linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support transfomers v5 (#994)

Browse Source 
>[!IMPORTANT]
>Do not merge this PR before all issues are resolved! 
>
>Testing with three versions: `4.52.0` , `4.57.6` and the latest stable
version

>[!NOTE]
>nvi-ci is split into correctness test ci and convergence test ci to
speed up testing in this PR
>, and more jobs for testing bc with transformers v4 (4.49.0 and 4.57.6)
>
>Whether keeping this change or not is yet to be discussed

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This is a dev branch for aggregating PRs related to transformers v5
changes.

<!---
## 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

---------

Signed-off-by: Tcc0403 <[email protected]>
Co-authored-by: Yunwei Dai <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
Co-authored-by: Yuki Uehara <[email protected]>
Co-authored-by: Chun-Mao Lai <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
Co-authored-by: Michael Lai <[email protected]>
This commit is contained in:
Tcc0403
2026-02-12 00:01:06 +00:00
committed by GitHub
co-authored by Yunwei Dai Vaibhav Jindal Yuki Uehara Chun-Mao Lai Vaibhav Jindal Michael Lai
parent 7d82f73d1e
commit 175cdf0e37
36 changed files with 1140 additions and 1857 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/nvi-ci.yml
+109 -5
View File
@@ -46,7 +46,7 @@ jobs:
- name: Run checkstyle
run: make checkstyle
tests:
correctness-tests:
runs-on: ubuntu-latest
needs: [checkstyle]
env:
@@ -69,15 +69,14 @@ jobs:
- name: Run tests
run: |
modal run dev.modal.tests
modal run dev.modal.tests::liger_correctness_tests
tests-bwd:
convergence-tests:
runs-on: ubuntu-latest
needs: [checkstyle]
env:
MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
REBUILD_IMAGE: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
steps:
- name: Checkout code
@@ -95,4 +94,109 @@ jobs:
- name: Run tests
run: |
modal run dev.modal.tests_bwd
modal run dev.modal.tests::liger_convergence_tests
correctness-tests-with-transformers-4-52-0:
runs-on: ubuntu-latest
needs: [checkstyle]
env:
MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
steps:
- name: Checkout code
uses: actions/checkout@v6
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install modal
- name: Run tests
run: |
modal run dev.modal.tests::liger_oldest_v4_correctness_tests
convergence-tests-with-transformers-4-52-0:
runs-on: ubuntu-latest
needs: [checkstyle]
env:
MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
steps:
- name: Checkout code
uses: actions/checkout@v6
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install modal
- name: Run tests
run: |
modal run dev.modal.tests::liger_oldest_v4_convergence_tests
correctness-tests-with-transformers-4-57-6:
runs-on: ubuntu-latest
needs: [checkstyle]
env:
MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
steps:
- name: Checkout code
uses: actions/checkout@v6
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install modal
- name: Run tests
run: |
modal run dev.modal.tests::liger_latest_v4_correctness_tests
convergence-tests-with-transformers-4-57-6:
runs-on: ubuntu-latest
needs: [checkstyle]
env:
MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
steps:
- name: Checkout code
uses: actions/checkout@v6
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install modal
- name: Run tests
run: |
modal run dev.modal.tests::liger_latest_v4_convergence_tests
benchmark/scripts/benchmark_llama4_rope.py
-4
View File
@@ -40,8 +40,6 @@ def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRu
num_key_value_heads=num_kv_heads,
head_dim=head_dim,
max_position_embeddings=seq_len,
rope_theta=10000.0,
rope_scaling=None, # Use default rope type
)
rotary_emb = transformers_version_dispatch(
@@ -134,8 +132,6 @@ def bench_memory_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkR
num_key_value_heads=num_kv_heads,
head_dim=head_dim,
max_position_embeddings=seq_len,
rope_theta=10000.0,
rope_scaling=None, # Use default rope type
)
rotary_emb = transformers_version_dispatch(
dev/modal/tests.py
+97 -1
View File
@@ -6,6 +6,8 @@ ROOT_PATH = Path(__file__).parent.parent.parent
REMOTE_ROOT_PATH = "/root/liger-kernel"
PYTHON_VERSION = "3.12"
OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION = "4.52.0"
image = modal.Image.debian_slim(python_version=PYTHON_VERSION).pip_install("uv")
app = modal.App("liger_tests", image=image)
@@ -15,7 +17,7 @@ repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
@app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_tests():
def liger_correctness_tests():
import subprocess
subprocess.run(
@@ -25,4 +27,98 @@ def liger_tests():
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
@app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_convergence_tests():
import subprocess
subprocess.run(
["uv pip install -e '.[dev]' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
oldest_v4_app = modal.App("liger_oldest_v4_tests", image=image) # 4.52.0
@oldest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_oldest_v4_correctness_tests():
import subprocess
subprocess.run(
["uv pip install -e '.[dev]' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(
[f"uv pip install 'transformers=={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
@oldest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_oldest_v4_convergence_tests():
import subprocess
subprocess.run(
["uv pip install -e '.[dev]' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(
[f"uv pip install 'transformers=={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
latest_v4_app = modal.App("liger_latest_v4_tests", image=image) # 4.57.6
@latest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_latest_v4_correctness_tests():
import subprocess
subprocess.run(
["uv pip install -e '.[dev]' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(
[f"uv pip install 'transformers>={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}, <5.0.0' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
@latest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_latest_v4_convergence_tests():
import subprocess
subprocess.run(
["uv pip install -e '.[dev]' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(
[f"uv pip install 'transformers>={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}, <5.0.0' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
dev/modal/tests_bwd.py
-35
View File
@@ -1,35 +0,0 @@
from pathlib import Path
import modal
ROOT_PATH = Path(__file__).parent.parent.parent
REMOTE_ROOT_PATH = "/root/liger-kernel"
PYTHON_VERSION = "3.12"
image = modal.Image.debian_slim(python_version=PYTHON_VERSION).pip_install("uv")
app = modal.App("liger_tests_bwd", image=image)
# mount: add local files to the remote container
repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
@app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_bwd_tests():
import subprocess
subprocess.run(
["uv pip install -e '.[dev]' --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
# force install transformers==4.52.0
subprocess.run(
["uv pip install transformers==4.52.0 --system"],
check=True,
shell=True,
cwd=REMOTE_ROOT_PATH,
)
subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
setup.py
+1 -1
View File
@@ -33,7 +33,7 @@ def get_optional_dependencies():
"""Get optional dependency groups."""
return {
"dev": [
"transformers>=4.52.0, <5.0.0",
"transformers>=4.52.0",
"matplotlib>=3.7.2",
"ruff>=0.12.0",
"pytest>=7.1.2",
src/liger_kernel/transformers/__init__.py
+1
View File
@@ -21,6 +21,7 @@ from liger_kernel.transformers.rope import liger_rotary_pos_emb # noqa: F401
from liger_kernel.transformers.softmax import LigerSoftmax # noqa: F401
from liger_kernel.transformers.sparsemax import LigerSparsemax # noqa: F401
from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP # noqa: F401
from liger_kernel.transformers.swiglu import LigerExperts # noqa: F401
from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP # noqa: F401
from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP # noqa: F401
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP # noqa: F401
src/liger_kernel/transformers/model/gemma.py
-121
View File
@@ -5,134 +5,13 @@ from typing import Union
import torch
from torch.nn import CrossEntropyLoss
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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_deprecated(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = 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,
skip_logits: Optional[bool] = None,
) -> Union[Tuple, CausalLMOutputWithPast]:
r"""
copy paste transformers.models.gemma.modeling_gemma causalLM with loss replaced with liger fused 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]`.
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, GemmaForCausalLM
>>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
>>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
>>> prompt = "What is your favorite condiment?"
>>> 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]
"What is your favorite condiment?"
```"""
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,
)
hidden_states = outputs[0]
loss = None
logits = None
if skip_logits and labels is None:
raise ValueError("skip_logits is True, but labels is None")
if skip_logits is None:
# By default, if in training mode, don't materialize logits
skip_logits = self.training and labels is not None
if skip_logits:
shift_hidden_states = hidden_states[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# flatten
shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
shift_labels = shift_labels.view(-1)
lce = LigerFusedLinearCrossEntropyLoss()
loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
else:
logits = self.lm_head(hidden_states)
if labels is not None:
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Ensure tensors are on the same device
shift_labels = shift_labels.to(shift_logits.device)
loss_fct = CrossEntropyLoss()
loss = loss_fct(shift_logits, shift_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/gemma2.py
-129
View File
@@ -6,12 +6,8 @@ from typing import Union
import torch
from torch.nn import CrossEntropyLoss
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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
@@ -19,131 +15,6 @@ from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWi
logger = logging.getLogger(__name__)
def lce_forward_deprecated(
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,
skip_logits: Optional[bool] = None,
**kwargs,
) -> Union[Tuple, CausalLMOutputWithPast]:
r"""
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]`.
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, GemmaForCausalLM
>>> model = GemmaForCausalLM.from_pretrained("google/gemma-2-9b")
>>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
>>> prompt = "What is your favorite condiment?"
>>> 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]
"What is your favorite condiment?"
```"""
if self.training and self.config._attn_implementation != "eager":
logger.warning_once(
"It is strongly recommended to train Gemma2 models with the `eager` attention implementation "
f"instead of `{self.config._attn_implementation}`. Use `eager` with `AutoModelForCausalLM.from_pretrained('<path-to-checkpoint>', attn_implementation='eager')`."
)
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]
loss = None
logits = None
if skip_logits and labels is None:
raise ValueError("skip_logits is True, but labels is None")
if skip_logits is None:
# By default, if in training mode, don't materialize logits
skip_logits = self.training and labels is not None
if skip_logits:
shift_hidden_states = hidden_states[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# flatten
shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
shift_labels = shift_labels.view(-1)
lce = LigerFusedLinearCrossEntropyLoss(softcap=self.config.final_logit_softcapping)
loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
else:
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
logits = self.lm_head(hidden_states)
if self.config.final_logit_softcapping is not None:
logits = logits / self.config.final_logit_softcapping
logits = torch.tanh(logits)
logits = logits * self.config.final_logit_softcapping
loss = None
if labels is not None:
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/glm4.py
-3
View File
@@ -5,14 +5,11 @@ from typing import Union
import torch
from transformers.utils.deprecation import deprecate_kwarg
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
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/glm4v.py
-3
View File
@@ -5,14 +5,11 @@ from typing import Union
import torch
from transformers.utils.deprecation import deprecate_kwarg
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
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/glm4v_moe.py
-3
View File
@@ -4,14 +4,11 @@ from typing import Union
import torch
from transformers.utils.deprecation import deprecate_kwarg
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 LigerGlm4vMoeCausalLMOutputWithPast
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/llama.py
-127
View File
@@ -5,15 +5,10 @@ from typing import Tuple
from typing import Union
import torch
import torch.nn.functional as F
from torch.distributed.fsdp import FullyShardedDataParallel
from torch.nn import CrossEntropyLoss
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fsdp import _FSDPForwardRedirection
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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
@@ -26,128 +21,6 @@ if PEFT_AVAILABLE:
from peft.utils.other import ModulesToSaveWrapper
def lce_forward_deprecated(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Union["Cache", List[torch.FloatTensor]]] = 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,
skip_logits: Optional[bool] = None,
) -> Union[Tuple, CausalLMOutputWithPast]:
r"""
Copy paste llama 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]`.
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, LlamaForCausalLM
>>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
>>> 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,
)
hidden_states = outputs[0]
loss = None
logits = None
# if in training mode, don't materialize logits
if skip_logits and labels is None:
raise ValueError("skip_logits is True, but labels is None")
if skip_logits is None:
# By default, if in training mode, don't materialize logits
skip_logits = self.training and labels is not None
if skip_logits:
shift_hidden_states = hidden_states[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# flatten tokens
shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
shift_labels = shift_labels.view(-1)
lce = LigerFusedLinearCrossEntropyLoss()
loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
else:
if self.config.pretraining_tp > 1:
lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
logits = torch.cat(logits, dim=-1)
else:
logits = self.lm_head(hidden_states)
if labels is not None:
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/llava.py
-187
View File
@@ -5,198 +5,11 @@ from typing import Union
import torch
from torch.nn import CrossEntropyLoss
from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast
from transformers.utils import is_torchdynamo_compiling
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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 LigerLlavaCausalLMOutputWithPast
def lce_forward_deprecated(
self,
input_ids: torch.LongTensor = None,
pixel_values: torch.FloatTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
vision_feature_layer: Optional[int] = None,
vision_feature_select_strategy: Optional[str] = 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,
image_sizes: torch.Tensor = None,
skip_logits: Optional[bool] = None,
**lm_kwargs,
) -> Union[Tuple, LlavaCausalLMOutputWithPast]:
r"""
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 PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, LlavaForConditionalGeneration
>>> model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf")
>>> processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
>>> prompt = "USER: <image>\nWhat's the content of the image? ASSISTANT:"
>>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> inputs = processor(images=image, text=prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(**inputs, max_new_tokens=15)
>>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"USER: \nWhat's the content of the image? ASSISTANT: The image features a busy city street with a stop sign prominently displayed"
```"""
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
vision_feature_layer = (
vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
)
vision_feature_select_strategy = (
vision_feature_select_strategy
if vision_feature_select_strategy is not None
else self.config.vision_feature_select_strategy
)
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
if pixel_values is not None and inputs_embeds is not None:
raise ValueError(
"You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
)
if inputs_embeds is None:
inputs_embeds = self.get_input_embeddings()(input_ids)
if pixel_values is not None:
image_features = self.get_image_features(
pixel_values=pixel_values,
vision_feature_layer=vision_feature_layer,
vision_feature_select_strategy=vision_feature_select_strategy,
image_sizes=image_sizes,
)
special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel():
n_image_tokens = (input_ids == self.config.image_token_index).sum()
n_image_features = image_features.shape[0] * image_features.shape[1]
raise ValueError(
f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
)
image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
outputs = self.language_model.model(
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,
logits_to_keep=logits_to_keep,
**lm_kwargs,
)
hidden_states = outputs[0]
loss = None
logits = None
# Overwrite skip_logits, since llava never materializes logits
skip_logits = labels is not None
if skip_logits:
# Shift so that tokens < n predict n
if attention_mask is not None:
# we use the input attention mask to shift the logits and labels, because it is 2D.
# we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
shift_attention_mask = attention_mask[:, -(hidden_states.shape[1] - 1) :].to(hidden_states.device)
shift_hidden_states = hidden_states[..., :-1, :][
shift_attention_mask.to(hidden_states.device) != 0
].contiguous()
shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
else:
shift_hidden_states = hidden_states[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
lce = LigerFusedLinearCrossEntropyLoss()
loss = lce(
self.language_model.lm_head.weight,
shift_hidden_states.view(-1, shift_hidden_states.size(-1)),
shift_labels.view(-1).to(shift_hidden_states.device),
)
else:
logits = self.language_model.lm_head(hidden_states)
if labels is not None:
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
shift_logits = logits[..., :-1, :]
shift_labels = labels[..., 1:]
if attention_mask is not None:
# we use the input attention mask to shift the logits and labels, because it is 2D.
# we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
else:
shift_logits = shift_logits.contiguous()
shift_labels = shift_labels.contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
flat_labels = shift_labels.view(-1).to(shift_logits.device)
loss = loss_fct(flat_logits, flat_labels)
if not return_dict:
# NOTE: This part has not been tested.
output = outputs[1:]
return (loss,) + output if loss is not None else output
return LlavaCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
image_hidden_states=image_features if pixel_values is not None else None,
)
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/mistral.py
-2
View File
@@ -6,14 +6,12 @@ from typing import Union
import torch
from transformers.cache_utils import Cache
from transformers.utils.deprecation import deprecate_kwarg
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
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/mixtral.py
-130
View File
@@ -5,143 +5,13 @@ from typing import Union
import torch
from torch.nn import CrossEntropyLoss
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
from transformers.models.mixtral.modeling_mixtral import load_balancing_loss_func
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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 LigerMoeCausalLMOutputWithPast
def lce_forward_deprecated(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = 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,
output_router_logits: Optional[bool] = None,
return_dict: Optional[bool] = None,
cache_position: Optional[torch.LongTensor] = None,
) -> Union[Tuple, MoeCausalLMOutputWithPast]:
r"""
Copy paste Mixtral's forward from transformers v4.44.2 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]`.
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, MixtralForCausalLM
>>> model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
>>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
>>> 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_router_logits = (
output_router_logits if output_router_logits is not None else self.config.output_router_logits
)
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,
output_router_logits=output_router_logits,
return_dict=return_dict,
cache_position=cache_position,
)
hidden_states = outputs[0]
logits = self.lm_head(hidden_states)
loss = None
if self.training and (labels is not None):
shift_hidden_states = hidden_states[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
shift_labels = shift_labels.view(-1)
lce = LigerFusedLinearCrossEntropyLoss()
loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
elif labels is not None:
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.weight, shift_labels)
aux_loss = None
if output_router_logits:
aux_loss = load_balancing_loss_func(
outputs.router_logits if return_dict else outputs[-1],
self.num_experts,
self.num_experts_per_tok,
attention_mask,
)
if labels is not None:
loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
if not return_dict:
output = (logits,) + outputs[1:]
if output_router_logits:
output = (aux_loss,) + output
return (loss,) + output if loss is not None else output
return MoeCausalLMOutputWithPast(
loss=loss,
aux_loss=aux_loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
router_logits=outputs.router_logits,
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
# Ignore copy
def lce_forward(
self,
src/liger_kernel/transformers/model/mllama.py
+1 -124
View File
@@ -5,133 +5,13 @@ from typing import Union
import torch
from torch.nn import CrossEntropyLoss
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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_deprecated(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
cross_attention_states: Optional[torch.LongTensor] = None,
cross_attention_mask: Optional[torch.LongTensor] = None,
full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = 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,
num_logits_to_keep: int = 0,
) -> Union[Tuple, CausalLMOutputWithPast]:
r"""
Copy paste mllama 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]`.
num_logits_to_keep (`int`, *optional*):
Calculate logits for the last `num_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.
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, MllamaForCausalLM
>>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
>>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
>>> prompt = "If I had to write a haiku, it would be:"
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6)
>>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
>>> print(result)
If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful.
I love the idea of snowflakes gently falling, each one
```
"""
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,
cross_attention_states=cross_attention_states,
attention_mask=attention_mask,
position_ids=position_ids,
cross_attention_mask=cross_attention_mask,
full_text_row_masked_out_mask=full_text_row_masked_out_mask,
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,
)
hidden_states = outputs[0]
loss = None
logits = None
if self.training and (labels is not None):
kept_hidden_states = hidden_states[:, -num_logits_to_keep:, :]
shift_hidden_states = kept_hidden_states[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# flatten tokens
shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
shift_labels = shift_labels.view(-1)
lce = LigerFusedLinearCrossEntropyLoss()
loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
else:
logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float()
if labels is not None:
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
@@ -192,9 +72,6 @@ def lce_forward(
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
# Filter out accum_dtype from kwargs for model call as MllamaTextModel doesn't accept it in transformers 4.49.0
# but preserve it for loss function calls
model_kwargs = {k: v for k, v in kwargs.items() if k != "accum_dtype"}
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
outputs = self.model(
input_ids=input_ids,
@@ -210,7 +87,7 @@ def lce_forward(
output_hidden_states=output_hidden_states,
return_dict=return_dict,
cache_position=cache_position,
**model_kwargs,
**kwargs,
)
hidden_states = outputs[0]
src/liger_kernel/transformers/model/olmo2.py
-3
View File
@@ -5,14 +5,11 @@ from typing import Union
import torch
from transformers.utils.deprecation import deprecate_kwarg
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
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/olmo3.py
-2
View File
@@ -6,14 +6,12 @@ from typing import Union
import torch
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
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
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/paligemma.py
-186
View File
@@ -7,12 +7,9 @@ import torch
from torch.nn import CrossEntropyLoss
from transformers.cache_utils import Cache
from transformers.models.paligemma.modeling_paligemma import PaliGemmaCausalLMOutputWithPast
from transformers.utils import is_torchdynamo_compiling
from transformers.utils import logging
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
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 LigerPaliGemmaCausalLMOutputWithPast
@@ -20,189 +17,6 @@ from liger_kernel.transformers.model.output_classes import LigerPaliGemmaCausalL
logger = logging.get_logger(__name__)
def lce_forward_deprecated(
self,
input_ids: torch.LongTensor = None,
pixel_values: torch.FloatTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None,
token_type_ids: Optional[torch.LongTensor] = None,
cache_position: 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,
) -> Union[Tuple, PaliGemmaCausalLMOutputWithPast]:
r"""
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]`.
Returns:
Example:
```python
>>> from PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
>>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/PaliGemma-test-224px-hf")
>>> processor = AutoProcessor.from_pretrained("google/PaliGemma-test-224px-hf")
>>> prompt = "answer en Where is the cow standing?"
>>> url = "https://huggingface.co/gv-hf/PaliGemma-test-224px-hf/resolve/main/cow_beach_1.png"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> inputs = processor(text=prompt, images=image, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(**inputs, max_length=30)
>>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"answer en Where is the cow standing?\nbeach"
```"""
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError(
"You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
)
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
# the attention mask is turned 4d after, we keep track of the original one
input_attention_mask = attention_mask
if inputs_embeds is None:
# 1. Extra the input embeddings
inputs_embeds = self.get_input_embeddings()(input_ids)
# 2. Merge text and images
if pixel_values is not None and input_ids.shape[1] != 1:
image_outputs = self.vision_tower(pixel_values.to(inputs_embeds.dtype))
selected_image_feature = image_outputs.last_hidden_state
image_features = self.multi_modal_projector(selected_image_feature)
if cache_position is None:
cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)
inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features(
image_features, inputs_embeds, input_ids, attention_mask, labels, token_type_ids, cache_position
)
else:
# In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
# generation with cache
if past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1:
# Retrieve the first layer to inspect the logits and mask out the hidden states
# that are set to 0
# TODO @molbap this will only work for dynamic cache.
first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]
# Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0)
# Get the target length
target_seqlen = cache_position[-1] + 1
extended_attention_mask = torch.ones(
(attention_mask.shape[0], target_seqlen - attention_mask.shape[1] + 1),
dtype=attention_mask.dtype,
device=attention_mask.device,
)
# Filter out only the tokens that can be un-attended, this can happen
# if one uses PaliGemma+ Fused modules where the cache on the
# first iteration is already big enough, or if one passes custom cache
valid_indices = non_attended_tokens < extended_attention_mask.size(-1)
new_batch_index = batch_index[valid_indices]
new_non_attended_tokens = non_attended_tokens[valid_indices]
# Zero-out the places where we don't need to attend
extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0
attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1)
position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
attention_mask = attention_mask.to(inputs_embeds.dtype)
outputs = self.language_model.model(
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,
)
hidden_states = outputs[0]
loss = None
logits = None
if self.training and (labels is not None):
shift_hidden_states = hidden_states[..., :-1, :]
shift_labels = labels[..., 1:]
hidden_device = shift_hidden_states.device
if attention_mask is not None:
# we use the input attention mask to shift the hidden_states and labels, because it is 2D.
# we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
shift_attention_mask = attention_mask[:, -shift_hidden_states.shape[1] :].to(hidden_device)
shift_hidden_states = shift_hidden_states[shift_attention_mask.to(hidden_device) != 0].contiguous()
shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
else:
shift_hidden_states = shift_hidden_states.contiguous()
shift_labels = shift_labels.contiguous()
# Flatten hidden state
shift_hidden_states = shift_hidden_states.view(-1, self.config.text_config.hidden_size)
shift_labels = shift_labels.view(-1).to(hidden_device)
lce = LigerFusedLinearCrossEntropyLoss()
loss = lce(self.language_model.lm_head.weight, shift_hidden_states, shift_labels)
else:
logits = self.language_model.lm_head(hidden_states)
if labels is not None:
shift_logits = logits[..., :-1, :]
shift_labels = labels[..., 1:]
if input_attention_mask is not None:
# we use the input attention mask to shift the logits and labels, because it is 2D.
shift_attention_mask = input_attention_mask[..., 1:]
shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
else:
shift_logits = shift_logits.contiguous()
shift_labels = shift_labels.contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
flat_logits = shift_logits.view(-1, self.config.vocab_size)
flat_labels = shift_labels.view(-1).to(shift_logits.device)
loss = loss_fct(flat_logits, flat_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return PaliGemmaCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/qwen2.py
-2
View File
@@ -7,7 +7,6 @@ import torch
from torch.nn import CrossEntropyLoss
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
@@ -130,7 +129,6 @@ def lce_forward_deprecated(
)
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/model/qwen2_5_vl.py
+20 -2
View File
@@ -5,12 +5,30 @@ from typing import Union
import torch
from packaging import version
from transformers import __version__ as transformers_version
from transformers.utils import can_return_tuple
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 LigerQwen2_5_VLCausalLMOutputWithPast
_TRANSFORMERS_V5_OR_LATER = version.parse(transformers_version) >= version.parse("5.0.0")
def _get_hidden_size(config) -> int:
"""Get hidden_size from Qwen2.5VLConfig in a version-aware manner."""
if _TRANSFORMERS_V5_OR_LATER:
return config.text_config.hidden_size
return config.hidden_size
def _get_vocab_size(config) -> int:
"""Get vocab_size from Qwen2.5VLConfig in a version-aware manner."""
if _TRANSFORMERS_V5_OR_LATER:
return config.text_config.vocab_size
return config.vocab_size
@can_return_tuple
def lce_forward(
@@ -129,7 +147,7 @@ def lce_forward(
lm_head_weight=self.lm_head.weight,
labels=labels,
shift_labels=shift_labels,
hidden_size=self.config.hidden_size,
hidden_size=_get_hidden_size(self.config),
**kwargs,
)
loss, _, token_accuracy = unpack_cross_entropy_result(result)
@@ -142,7 +160,7 @@ def lce_forward(
logits=logits,
labels=labels,
shift_labels=shift_labels,
vocab_size=self.config.vocab_size,
vocab_size=_get_vocab_size(self.config),
)
if not return_dict:
src/liger_kernel/transformers/model/qwen2_vl.py
+20 -2
View File
@@ -5,12 +5,30 @@ from typing import Union
import torch
from packaging import version
from transformers import __version__ as transformers_version
from transformers.utils import can_return_tuple
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 LigerQwen2VLCausalLMOutputWithPast
_TRANSFORMERS_V5_OR_LATER = version.parse(transformers_version) >= version.parse("5.0.0")
def _get_hidden_size(config) -> int:
"""Get hidden_size from Qwen2VLConfig in a version-aware manner."""
if _TRANSFORMERS_V5_OR_LATER:
return config.text_config.hidden_size
return config.hidden_size
def _get_vocab_size(config) -> int:
"""Get vocab_size from Qwen2VLConfig in a version-aware manner."""
if _TRANSFORMERS_V5_OR_LATER:
return config.text_config.vocab_size
return config.vocab_size
@can_return_tuple
def lce_forward(
@@ -125,7 +143,7 @@ def lce_forward(
lm_head_weight=self.lm_head.weight,
labels=labels,
shift_labels=shift_labels,
hidden_size=self.config.hidden_size,
hidden_size=_get_hidden_size(self.config),
**kwargs,
)
loss, _, token_accuracy = unpack_cross_entropy_result(result)
@@ -138,7 +156,7 @@ def lce_forward(
logits=logits,
labels=labels,
shift_labels=shift_labels,
vocab_size=self.config.vocab_size,
vocab_size=_get_vocab_size(self.config),
)
if not return_dict:
src/liger_kernel/transformers/model/smollm3.py
-2
View File
@@ -7,7 +7,6 @@ from typing import Union
import torch
from torch.distributed.fsdp import FullyShardedDataParallel
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.fsdp import _FSDPForwardRedirection
from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
@@ -22,7 +21,6 @@ if PEFT_AVAILABLE:
from peft.utils.other import ModulesToSaveWrapper
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
src/liger_kernel/transformers/monkey_patch.py
+100 -172
View File
@@ -17,26 +17,21 @@ from liger_kernel.transformers.geglu import LigerGEGLUMLP
from liger_kernel.transformers.layer_norm import LigerLayerNorm
from liger_kernel.transformers.model.falcon_h1 import lce_forward as falcon_h1_lce_forward
from liger_kernel.transformers.model.gemma import lce_forward as gemma_lce_forward
from liger_kernel.transformers.model.gemma import lce_forward_deprecated as gemma_lce_forward_deprecated
from liger_kernel.transformers.model.gemma2 import lce_forward as gemma2_lce_forward
from liger_kernel.transformers.model.gemma2 import lce_forward_deprecated as gemma2_lce_forward_deprected
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.llama import lce_forward_deprecated as llama_lce_forward_deprecated
from liger_kernel.transformers.model.llava import lce_forward as llava_lce_forward
from liger_kernel.transformers.model.llava import lce_forward_deprecated as llava_lce_forward_deprecated
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.mixtral import lce_forward_deprecated as mixtral_lce_forward_deprecated
from liger_kernel.transformers.model.phi3 import lce_forward as phi3_lce_forward
from liger_kernel.transformers.model.qwen2 import lce_forward as qwen2_lce_forward
from liger_kernel.transformers.model.qwen2 import lce_forward_deprecated as qwen2_lce_forward_deprecated
from liger_kernel.transformers.model.smollm3 import lce_forward as smollm3_lce_forward
from liger_kernel.transformers.qwen2vl_mrope import liger_multimodal_rotary_pos_emb
from liger_kernel.transformers.rms_norm import LigerRMSNorm
from liger_kernel.transformers.rope import liger_rotary_pos_emb
from liger_kernel.transformers.rope import liger_rotary_pos_emb_vision
from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP
from liger_kernel.transformers.swiglu import LigerExperts
from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
@@ -50,8 +45,15 @@ except ImportError:
transformer_version = version.parse(transformers.__version__)
logger = logging.getLogger(__name__)
SUPPORTED_TRANSFORMER_VERSION = "4.46.1"
TRANSFORMER_DEPRECATION_WARNING = "Support for transformers versions < 4.46.1 will soon be discontinued due to issues with incorrect gradient accumulation. \n Please consider upgrading to avoid potential issues. See details: https://github.com/huggingface/transformers/pull/34191"
MIN_SUPPORTED_TRANSFORMERS_VERSION = version.parse("4.52.0")
if transformer_version < MIN_SUPPORTED_TRANSFORMERS_VERSION:
raise ImportError(
f"liger-kernel requires transformers >= {MIN_SUPPORTED_TRANSFORMERS_VERSION}, got {transformers.__version__}. "
"Please install an older version of liger-kernel that is compatible with your transformers version."
)
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
def _bind_method_to_module(module, method_name: str, new_method: Callable):
@@ -183,13 +185,9 @@ def apply_liger_kernel_to_granite(
modeling_granite.apply_rotary_pos_emb = liger_rotary_pos_emb
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_granite.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
raise NotImplementedError("LigerFusedLinearCrossEntropy is not available for Granite models.")
@@ -255,26 +253,15 @@ def apply_liger_kernel_to_llama(
modeling_llama.LlamaMLP = LigerSwiGLUMLP
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_llama.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(llama_lce_forward, model)
else:
modeling_llama.LlamaForCausalLM.forward = llama_lce_forward
else: # if version < 4.46.1
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(llama_lce_forward_deprecated, model)
else:
modeling_llama.LlamaForCausalLM.forward = llama_lce_forward_deprecated
if model is not None:
model.forward = MethodType(llama_lce_forward, model)
else:
modeling_llama.LlamaForCausalLM.forward = llama_lce_forward
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -333,13 +320,9 @@ def apply_liger_kernel_to_smollm3(
modeling_smollm3.SmolLM3MLP = LigerSwiGLUMLP
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_smollm3.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if model is not None:
@@ -396,23 +379,14 @@ def apply_liger_kernel_to_llava(
from transformers.models.llava import modeling_llava
if cross_entropy:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_llava.nn.CrossEntropyLoss = LigerCrossEntropyLoss
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if transformer_version >= version.parse("4.52.0"):
if model is not None:
model.forward = MethodType(llava_lce_forward, model)
else:
modeling_llava.LlavaForConditionalGeneration.forward = llava_lce_forward
elif transformer_version >= version.parse("4.49.0") and transformer_version < version.parse("4.52.0"):
if model is not None:
model.forward = MethodType(llava_lce_forward_deprecated, model)
else:
modeling_llava.LlavaForConditionalGeneration.forward = llava_lce_forward_deprecated
else: # if version < 4.49.0
logger.warning(
"The latest version of Liger does not support transformers < 4.49.0 for llava. Please downgrade your liger version or upgrade your transformer version."
)
if model is not None:
model.forward = MethodType(llava_lce_forward, model)
else:
modeling_llava.LlavaForConditionalGeneration.forward = llava_lce_forward
if model is not None:
text_model_name, vision_model_name = model.config.text_config.model_type, model.config.vision_config.model_type
@@ -579,7 +553,6 @@ def apply_liger_kernel_to_mllama(
from transformers.models.mllama.modeling_mllama import MllamaVisionModel
from liger_kernel.transformers.model.mllama import lce_forward as mllama_lce_forward
from liger_kernel.transformers.model.mllama import lce_forward_deprecated as mllama_lce_forward_deprecated
if rope:
modeling_mllama.apply_rotary_pos_emb = liger_rotary_pos_emb
@@ -590,25 +563,14 @@ def apply_liger_kernel_to_mllama(
if swiglu:
modeling_mllama.MllamaTextMLP = LigerSwiGLUMLP
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_mllama.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(mllama_lce_forward, model)
else:
modeling_mllama.MllamaForCausalLM.forward = mllama_lce_forward
else: # if version < 4.46.1
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(mllama_lce_forward_deprecated, model)
else:
modeling_mllama.MllamaForCausalLM.forward = mllama_lce_forward_deprecated
if model is not None:
model.forward = MethodType(mllama_lce_forward, model)
else:
modeling_mllama.MllamaForCausalLM.forward = mllama_lce_forward
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -694,16 +656,10 @@ def apply_liger_kernel_to_mistral(
if cross_entropy:
modeling_mistral.CrossEntropyLoss = LigerCrossEntropyLoss
if fused_linear_cross_entropy:
if transformer_version >= version.parse("4.49.0"):
if model is not None:
model.forward = MethodType(mistral_lce_forward, model)
else:
modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
if model is not None:
model.forward = MethodType(mistral_lce_forward, model)
else:
logger.warning(
"The latest version of Liger does not support transformers < 4.49.0 for llava. Please downgrade your liger version or upgrade your transformer version."
)
logger.warning("LigerFusedLinearCrossEntropy patch is not applied.")
modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
if swiglu:
modeling_mistral.MistralMLP = LigerSwiGLUMLP
@@ -762,28 +718,20 @@ def apply_liger_kernel_to_mixtral(
if rms_norm:
modeling_mixtral.MixtralRMSNorm = LigerRMSNorm
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_mixtral.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(mixtral_lce_forward, model)
else:
modeling_mixtral.MixtralForCausalLM.forward = mixtral_lce_forward
else: # if version < 4.46.1
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(mixtral_lce_forward_deprecated, model)
else:
modeling_mixtral.MixtralForCausalLM.forward = mixtral_lce_forward_deprecated
if model is not None:
model.forward = MethodType(mixtral_lce_forward, model)
else:
modeling_mixtral.MixtralForCausalLM.forward = mixtral_lce_forward
if swiglu:
modeling_mixtral.MixtralBlockSparseTop2MLP = LigerBlockSparseTop2MLP
if IS_TRANSFORMERS_V5_OR_LATER:
modeling_mixtral.MixtralExperts = LigerExperts
else:
modeling_mixtral.MixtralBlockSparseTop2MLP = LigerBlockSparseTop2MLP
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -797,8 +745,11 @@ def apply_liger_kernel_to_mixtral(
for decoder_layer in base_model.layers:
if swiglu:
for expert in decoder_layer.block_sparse_moe.experts:
_patch_swiglu_module(expert, LigerBlockSparseTop2MLP)
if IS_TRANSFORMERS_V5_OR_LATER:
_patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)
else:
for expert in decoder_layer.block_sparse_moe.experts:
_patch_swiglu_module(expert, LigerBlockSparseTop2MLP)
if rms_norm:
_patch_rms_norm_module(decoder_layer.input_layernorm)
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -844,27 +795,16 @@ def apply_liger_kernel_to_gemma(
if rms_norm:
modeling_gemma.GemmaRMSNorm = LigerRMSNormForGemma
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_gemma.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if geglu:
modeling_gemma.GemmaMLP = LigerGEGLUMLP
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(gemma_lce_forward, model)
else:
modeling_gemma.GemmaForCausalLM.forward = gemma_lce_forward
else: # if version < 4.46.1
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(gemma_lce_forward_deprecated, model)
else:
modeling_gemma.GemmaForCausalLM.forward = gemma_lce_forward_deprecated
if model is not None:
model.forward = MethodType(gemma_lce_forward, model)
else:
modeling_gemma.GemmaForCausalLM.forward = gemma_lce_forward
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -927,25 +867,14 @@ def apply_liger_kernel_to_gemma2(
# https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/gemma/modeling_gemma.py#L109
modeling_gemma2.Gemma2RMSNorm = LigerRMSNormForGemma2
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_gemma2.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(gemma2_lce_forward, model)
else:
modeling_gemma2.Gemma2ForCausalLM.forward = gemma2_lce_forward
if model is not None:
model.forward = MethodType(gemma2_lce_forward, model)
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(gemma2_lce_forward_deprected, model)
else:
modeling_gemma2.Gemma2ForCausalLM.forward = gemma2_lce_forward_deprected
modeling_gemma2.Gemma2ForCausalLM.forward = gemma2_lce_forward
if geglu:
modeling_gemma2.Gemma2MLP = LigerGEGLUMLP
@@ -1189,7 +1118,6 @@ def apply_liger_kernel_to_paligemma(
from transformers.models.siglip.modeling_siglip import SiglipVisionModel
from liger_kernel.transformers.model.paligemma import lce_forward
from liger_kernel.transformers.model.paligemma import lce_forward_deprecated
# The vision_tower is a SiglipVisionModel
if layer_norm and model is None:
@@ -1209,17 +1137,10 @@ def apply_liger_kernel_to_paligemma(
if cross_entropy:
modeling_paligemma.nn.CrossEntropyLoss = LigerCrossEntropyLoss
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(lce_forward, model)
else:
modeling_paligemma.PaliGemmaForConditionalGeneration.forward = lce_forward
else: # if version < 4.46.1
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(lce_forward_deprecated, model)
else:
modeling_paligemma.PaliGemmaForConditionalGeneration.forward = lce_forward_deprecated
if model is not None:
model.forward = MethodType(lce_forward, model)
else:
modeling_paligemma.PaliGemmaForConditionalGeneration.forward = lce_forward
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -1301,26 +1222,15 @@ def apply_liger_kernel_to_qwen2(
modeling_qwen2.Qwen2RMSNorm = LigerRMSNorm
if cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from transformers.loss.loss_utils import nn
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
else:
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
modeling_qwen2.CrossEntropyLoss = LigerCrossEntropyLoss
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
if model is not None:
model.forward = MethodType(qwen2_lce_forward, model)
else:
modeling_qwen2.Qwen2ForCausalLM.forward = qwen2_lce_forward
else: # if version < 4.46.1
logger.warning(TRANSFORMER_DEPRECATION_WARNING)
if model is not None:
model.forward = MethodType(qwen2_lce_forward_deprecated, model)
else:
modeling_qwen2.Qwen2ForCausalLM.forward = qwen2_lce_forward_deprecated
if model is not None:
model.forward = MethodType(qwen2_lce_forward, model)
else:
modeling_qwen2.Qwen2ForCausalLM.forward = qwen2_lce_forward
if swiglu:
modeling_qwen2.Qwen2MLP = LigerSwiGLUMLP
@@ -1439,7 +1349,10 @@ def apply_liger_kernel_to_qwen3_moe(
modeling_qwen3_moe.Qwen3MoeForCausalLM.forward = qwen3_lce_forward
if swiglu:
modeling_qwen3_moe.Qwen3MoeMLP = LigerQwen3MoeSwiGLUMLP
if IS_TRANSFORMERS_V5_OR_LATER:
modeling_qwen3_moe.Qwen3MoeExperts = LigerExperts
else:
modeling_qwen3_moe.Qwen3MoeMLP = LigerQwen3MoeSwiGLUMLP
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -1452,8 +1365,11 @@ def apply_liger_kernel_to_qwen3_moe(
_patch_rms_norm_module(base_model.norm)
for decoder_layer in base_model.layers:
if swiglu:
for mlp_expert in decoder_layer.mlp.experts:
_patch_swiglu_module(mlp_expert, LigerQwen3MoeSwiGLUMLP)
if IS_TRANSFORMERS_V5_OR_LATER:
_patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)
else:
for mlp_expert in decoder_layer.mlp.experts:
_patch_swiglu_module(mlp_expert, LigerQwen3MoeSwiGLUMLP)
if rms_norm:
_patch_rms_norm_module(decoder_layer.input_layernorm)
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -2671,8 +2587,11 @@ def apply_liger_kernel_to_qwen3_next(
else:
modeling_qwen3_next.Qwen3NextForCausalLM.forward = qwen3_next_lce_forward
if swiglu:
# Qwen3MoeMLP and Qwen3NextMLP are identical, hence we reuse LigerQwen3MoeSwiGLUMLP
modeling_qwen3_next.Qwen3NextMLP = LigerQwen3MoeSwiGLUMLP
if IS_TRANSFORMERS_V5_OR_LATER:
modeling_qwen3_next.Qwen3NextExperts = LigerExperts
else:
# Qwen3MoeMLP and Qwen3NextMLP are identical, hence we reuse LigerQwen3MoeSwiGLUMLP
modeling_qwen3_next.Qwen3NextMLP = LigerQwen3MoeSwiGLUMLP
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -2700,8 +2619,11 @@ def apply_liger_kernel_to_qwen3_next(
_patch_swiglu_module(decoder_layer.mlp.shared_expert, LigerQwen3MoeSwiGLUMLP)
experts = getattr(decoder_layer.mlp, "experts", None)
if experts is not None:
for expert in experts:
_patch_swiglu_module(expert, LigerQwen3MoeSwiGLUMLP)
if IS_TRANSFORMERS_V5_OR_LATER:
_patch_swiglu_module(experts, LigerExperts)
else:
for expert in experts:
_patch_swiglu_module(expert, LigerQwen3MoeSwiGLUMLP)
def apply_liger_kernel_to_hunyuan_v1_dense(
@@ -2801,7 +2723,10 @@ def apply_liger_kernel_to_hunyuan_v1_moe(
modeling_hunyuan_v1_moe.HunYuanMoEV1ForCausalLM.forward = hunyuan_v1_moe_lce_forward
if swiglu:
modeling_hunyuan_v1_moe.HunYuanMoEV1MLP = LigerHunyuanV1SwiGLUMLP
if IS_TRANSFORMERS_V5_OR_LATER:
modeling_hunyuan_v1_moe.HunYuanMoEV1Experts = LigerExperts
else:
modeling_hunyuan_v1_moe.HunYuanMoEV1MLP = LigerHunyuanV1SwiGLUMLP
if model is not None:
# The model instance already exists, so we need to additionally patch the
@@ -2814,8 +2739,11 @@ def apply_liger_kernel_to_hunyuan_v1_moe(
_patch_rms_norm_module(base_model.norm)
for decoder_layer in base_model.layers:
if swiglu:
for mlp_expert in decoder_layer.mlp.experts:
_patch_swiglu_module(mlp_expert, LigerHunyuanV1SwiGLUMLP)
if IS_TRANSFORMERS_V5_OR_LATER:
_patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)
else:
for mlp_expert in decoder_layer.mlp.experts:
_patch_swiglu_module(mlp_expert, LigerHunyuanV1SwiGLUMLP)
if rms_norm:
_patch_rms_norm_module(decoder_layer.input_layernorm)
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
src/liger_kernel/transformers/swiglu.py
+49
View File
@@ -1,3 +1,4 @@
import torch
import torch.nn as nn
from liger_kernel.ops import LigerSiLUMulFunction
@@ -36,6 +37,54 @@ class LigerBlockSparseTop2MLP(nn.Module):
return self.w2(LigerSiLUMulFunction.apply(self.w1(x), self.w3(x)))
class LigerExperts(nn.Module):
"""
Patch MixtralExperts for transformers v5 or later to use LigerSiLUMulFunction
https://github.com/huggingface/transformers/blob/393b4b3d28e29b4b05b19b4b7f3242a7fc893637/src/transformers/models/mixtral/modeling_mixtral.py#L63
"""
def __init__(self, config):
super().__init__()
if "num_experts" in config:
# qwen3_moe, qwen3_next uses num_experts
self.num_experts = config.num_experts
else:
self.num_experts = config.num_local_experts
if "moe_intermediate_size" in config:
# qwen3_moe, qwen3_next uses moe_intermediate_size
self.intermediate_dim = config.moe_intermediate_size
else:
self.intermediate_dim = config.intermediate_size
self.hidden_dim = config.hidden_size
self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
if config.hidden_act not in ["silu", "swish"]:
raise ValueError(f"Activation function {config.hidden_act} not supported.")
def forward(self, hidden_states, top_k_index, top_k_weights):
final_hidden_states = torch.zeros_like(hidden_states)
with torch.no_grad():
expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
expert_mask = expert_mask.permute(2, 1, 0)
expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
for expert_idx in expert_hit:
expert_idx = expert_idx[0]
if expert_idx == self.num_experts:
continue
top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
current_state = hidden_states[token_idx]
gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
current_hidden_states = LigerSiLUMulFunction.apply(gate, up)
current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
return final_hidden_states
class LigerPhi3SwiGLUMLP(nn.Module):
"""
Patch Phi3MLP to use LigerSiLUMulFunction
test/convergence/bf16/test_mini_models.py
+105 -136
View File
@@ -4,8 +4,10 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # Ensure deterministic behavi
import pytest
import torch
import transformers
from datasets import load_from_disk
from packaging import version
from torch.utils.data import DataLoader
from transformers.models.gemma import GemmaConfig
from transformers.models.gemma import GemmaForCausalLM
@@ -53,6 +55,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_next
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl_moe
from liger_kernel.transformers import apply_liger_kernel_to_smollm3
from liger_kernel.utils import infer_device
from test.utils import DEFAULT_DATASET_PATH
from test.utils import MiniModelConfig
from test.utils import assert_verbose_allclose
@@ -94,6 +97,8 @@ from test.utils import set_seed
from test.utils import simple_collate_fn
from test.utils import supports_bfloat16
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
try:
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM
@@ -307,7 +312,6 @@ try:
except ImportError:
EXAONE4_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -333,8 +337,6 @@ MINI_MODEL_SETUPS = {
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -362,7 +364,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -391,7 +392,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=None, # defaults to num_attention_heads
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=None,
tie_word_embeddings=False,
use_cache=True,
@@ -416,7 +416,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -441,7 +440,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -476,7 +474,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -504,7 +501,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -532,7 +528,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -561,8 +556,6 @@ if LLAMA4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=10000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -590,7 +583,6 @@ if QWEN3_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -616,8 +608,6 @@ if QWEN3_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -653,14 +643,6 @@ if GPT_OSS_AVAILABLE:
rms_norm_eps=1e-5,
use_cache=True,
tie_word_embeddings=False,
rope_parameters={
"rope_type": "yarn",
"factor": 8.0,
"beta_fast": 32.0,
"beta_slow": 1.0,
"truncate": False,
"original_max_position_embeddings": 4096,
},
attention_dropout=0.0,
num_local_experts=8, # Reduced from 32 for mini model
num_experts_per_tok=2, # Reduced from 4 for mini model
@@ -693,7 +675,6 @@ if GEMMA3_AVAILABLE:
bos_token_id=2,
eos_token_id=1,
tie_word_embeddings=True,
rope_theta=10000.0, # 1000000
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -721,18 +702,20 @@ if MLLAMA_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_scaling=dict(
factor=8.0,
high_freq_factor=4.0,
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_theta=500_000,
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
)
@@ -742,36 +725,38 @@ if QWEN2_VL_AVAILABLE:
liger_kernel_patch_revert_func=revert_liger_kernel_to_qwen2_vl,
model_class=Qwen2VLForConditionalGeneration,
mini_model_config=Qwen2VLConfig(
attention_dropout=0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
bos_token_id=1, # 151643
eos_token_id=2, # 151645
# In transformers v5, text-related parameters must be in text_config
text_config={
"attention_dropout": 0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
"bos_token_id": 1, # 151643
"eos_token_id": 2, # 151645
"hidden_act": "silu",
"hidden_size": 1536, # 8192
"initializer_range": 0.02,
"intermediate_size": 4864, # 29568
"max_position_embeddings": 32768,
"max_window_layers": 4, # 80
"num_attention_heads": 12, # 64
"num_hidden_layers": 4, # 80
"num_key_value_heads": 2, # 8
"rms_norm_eps": 1e-6, # 1e-5
**(
{"rope_parameters": {"mrope_section": [16, 24, 24]}} # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else {"rope_scaling": {"type": "mrope", "mrope_section": [16, 24, 24]}}
),
"sliding_window": 4096,
"tie_word_embeddings": False,
"use_cache": True,
"vocab_size": 32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
"use_sliding_window": False,
},
vision_start_token_id=32765, # vocab_size - 5
vision_end_token_id=32766, # vocab_size - 4
vision_token_id=32767, # vocab_size - 3
image_token_id=32768, # vocab_size - 2
video_token_id=32769, # vocab_size - 1
hidden_act="silu",
hidden_size=1536, # 8192
initializer_range=0.02,
intermediate_size=4864, # 29568
max_position_embeddings=32768,
max_window_layers=4, # 80
num_attention_heads=12, # 64
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
),
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
use_sliding_window=False,
vision_config={
"depth": 4, # 32
"embed_dim": 1280,
@@ -784,7 +769,6 @@ if QWEN2_VL_AVAILABLE:
"spatial_patch_size": 14,
"temporal_patch_size": 2,
},
attn_implementation="sdpa",
),
)
@@ -794,36 +778,38 @@ if QWEN2_5_VL_AVAILABLE:
liger_kernel_patch_revert_func=revert_liger_kernel_to_qwen2_5_vl,
model_class=Qwen2_5_VLForConditionalGeneration,
mini_model_config=Qwen2_5_VLConfig(
attention_dropout=0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
bos_token_id=1, # 151643
eos_token_id=2, # 151645
# In transformers v5, text-related parameters must be in text_config
text_config={
"attention_dropout": 0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
"bos_token_id": 1, # 151643
"eos_token_id": 2, # 151645
"hidden_act": "silu",
"hidden_size": 1536, # 8192
"initializer_range": 0.02,
"intermediate_size": 4864, # 29568
"max_position_embeddings": 32768,
"max_window_layers": 4, # 80
"num_attention_heads": 12, # 64
"num_hidden_layers": 4, # 80
"num_key_value_heads": 2, # 8
"rms_norm_eps": 1e-6, # 1e-5
**(
{"rope_parameters": {"mrope_section": [16, 24, 24]}} # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else {"rope_scaling": {"type": "mrope", "mrope_section": [16, 24, 24]}}
),
"sliding_window": 4096,
"tie_word_embeddings": False,
"use_cache": True,
"vocab_size": 32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
"use_sliding_window": False,
},
vision_start_token_id=32765, # vocab_size - 5
vision_end_token_id=32766, # vocab_size - 4
vision_token_id=32767, # vocab_size - 3
image_token_id=32768, # vocab_size - 2
video_token_id=32769, # vocab_size - 1
hidden_act="silu",
hidden_size=1536, # 8192
initializer_range=0.02,
intermediate_size=4864, # 29568
max_position_embeddings=32768,
max_window_layers=4, # 80
num_attention_heads=12, # 64
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
),
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
use_sliding_window=False,
vision_config={
"depth": 4, # 32
"hidden_act": "silu",
@@ -840,7 +826,6 @@ if QWEN2_5_VL_AVAILABLE:
"tokens_per_second": 2,
"temporal_patch_size": 2,
},
attn_implementation="sdpa",
),
)
@@ -870,13 +855,15 @@ if QWEN3_VL_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
use_cache=True,
vocab_size=32768,
pad_token_id=None,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
vision_config=dict(
depth=4,
@@ -923,11 +910,6 @@ if QWEN3_VL_MOE_AVAILABLE:
num_key_value_heads=2,
head_dim=128,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
use_cache=True,
vocab_size=32768,
decoder_sparse_step=1,
@@ -936,6 +918,13 @@ if QWEN3_VL_MOE_AVAILABLE:
num_experts=4,
tie_word_embeddings=False,
mlp_only_layers=[],
pad_token_id=None,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
).to_dict(),
vision_config=Qwen3VLMoeVisionConfig(
depth=4,
@@ -977,8 +966,6 @@ if GRANITE_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1010,8 +997,6 @@ if LLAVA_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
pretraining_tp=1,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
max_position_embeddings=4096, # llava-1.5-7b-hf
@@ -1069,8 +1054,6 @@ if OLMO2_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1098,8 +1081,6 @@ if OLMO3_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1128,8 +1109,6 @@ if GLM4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1165,8 +1144,6 @@ if GLM4V_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1182,13 +1159,14 @@ if GLM4V_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
"pad_token_id": None,
},
vision_config={
"depth": 4, # 32
@@ -1232,8 +1210,6 @@ if GLM4V_MOE_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1249,11 +1225,6 @@ if GLM4V_MOE_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
"attention_dropout": 0.0,
@@ -1266,6 +1237,11 @@ if GLM4V_MOE_AVAILABLE:
"topk_group": 1,
"first_k_dense_replace": 1,
"norm_topk_prob": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
},
vision_config={
"depth": 4, # 32
@@ -1303,8 +1279,6 @@ if SMOLLM3_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1396,8 +1370,6 @@ if QWEN3NEXT_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -1437,7 +1409,6 @@ if HUNYUAN_V1_AVAILABLE:
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
@@ -1468,8 +1439,6 @@ if HUNYUAN_V1_AVAILABLE:
eod_token_id=3,
sep_token_id=4,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
attention_dropout=0.0,
num_experts=2,
@@ -1496,11 +1465,11 @@ if EXAONE4_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=1000000.0,
tie_word_embeddings=True,
use_cache=True,
vocab_size=32000,
attn_implementation="sdpa",
pad_token_id=None,
),
)
@@ -1602,12 +1571,12 @@ def run_mini_model(
"model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
[
pytest.param(
"mini_llama4",
"mini_llama4", # llama4 requires slightly larger tolerances to pass this test after bug fix to llama4 in transformers v5.0.0
32,
1e-5,
torch.bfloat16,
1e-2,
5e-2,
4e-1,
1e-1,
1e-1,
1e-2,
@@ -1618,6 +1587,10 @@ def run_mini_model(
not LLAMA4_AVAILABLE,
reason="Llama not available in this version of transformers",
),
pytest.mark.skipif(
not IS_TRANSFORMERS_V5_OR_LATER,
reason="The `attention_bias` configuration of Llama4 is not set in Transformers v4",
),
],
),
pytest.param(
@@ -1728,14 +1701,14 @@ def run_mini_model(
),
# TODO(tcc): Investigate qwen3_moe on different machines.
# The loss diverges on ci test (A10G), but it never diverges on my local machine (3080).
# Qwen3_moe can pass float32 tests.
# Qwen3_moe can pass float32 tests. (mecoli1219): diverges on h100
pytest.param(
"mini_qwen3_moe",
32,
1e-5,
torch.bfloat16,
5e-2,
5e-2,
2e-1,
1e-1, # 1e-1
1e-1, # 1e-2
1e-2,
@@ -1830,9 +1803,9 @@ def run_mini_model(
1e-5,
torch.bfloat16,
5e-2,
5e-2,
1e-1,
1e-2,
1e-1,
5e-2,
1e-2,
1e-2,
marks=[
@@ -1869,10 +1842,6 @@ def run_mini_model(
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
version.parse(transformers.__version__) < version.parse("4.49.0"),
reason="Mistral not available in transformers<=4.49.0",
),
],
),
pytest.param(
test/convergence/bf16/test_mini_models_multimodal.py
+61 -40
View File
@@ -4,11 +4,12 @@ import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # Ensure deterministic behavior with CuBLAS
import pytest
import torch
import transformers
from datasets import load_dataset
from packaging import version
from torch.utils.data import DataLoader
from transformers import PreTrainedTokenizerFast
from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
from liger_kernel.transformers import apply_liger_kernel_to_gemma3
@@ -22,6 +23,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen2_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl_moe
from liger_kernel.transformers import apply_liger_kernel_to_smolvlm
from liger_kernel.utils import infer_device
from test.utils import FAKE_CONFIGS_PATH
from test.utils import UNTOKENIZED_DATASET_PATH
from test.utils import MiniModelConfig
@@ -49,12 +51,23 @@ from test.utils import set_seed
from test.utils import supports_bfloat16
from test.utils import train_bpe_tokenizer
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.gemma.tokenization_gemma import GemmaTokenizer
else:
from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast as GemmaTokenizer
try:
# Qwen2-VL is only available in transformers>=4.52.4
import transformers
from packaging import version
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig
from transformers.models.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLForConditionalGeneration
@@ -70,7 +83,11 @@ try:
import transformers
from packaging import version
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
@@ -82,7 +99,10 @@ except ImportError:
QWEN2_5_VL_AVAILABLE = False
try:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
@@ -138,7 +158,6 @@ try:
from packaging import version
from transformers.models.gemma.configuration_gemma import GemmaConfig
from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
from transformers.models.gemma2.configuration_gemma2 import Gemma2Config
from transformers.models.paligemma.configuration_paligemma import PaliGemmaConfig
from transformers.models.paligemma.modeling_paligemma import PaliGemmaForConditionalGeneration
@@ -146,7 +165,7 @@ try:
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
from transformers.models.siglip.image_processing_siglip import SiglipImageProcessor
PALIGEMMA_AVAILABLE = version.parse(transformers.__version__) >= version.parse("4.46.0")
PALIGEMMA_AVAILABLE = True
except ImportError:
PALIGEMMA_AVAILABLE = False
@@ -177,7 +196,6 @@ except ImportError:
LLAMA4_AVAILABLE = False
try:
# InternVL is only available in transformers>=4.52.1
from transformers.models.got_ocr2.image_processing_got_ocr2_fast import GotOcr2ImageProcessorFast
from transformers.models.internvl.configuration_internvl import InternVLConfig
from transformers.models.internvl.modeling_internvl import InternVLForConditionalGeneration
@@ -185,13 +203,14 @@ try:
from transformers.models.internvl.video_processing_internvl import InternVLVideoProcessor
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
INTERNVL_AVAILABLE = True
# Input fp32 with bf16 CNN-based models in InternVL is only working in transformers>=4.56.0
INTERNVL_AVAILABLE = version.parse(transformers.__version__) >= version.parse("4.56.0")
except ImportError:
INTERNVL_AVAILABLE = False
try:
# SmolVLM2 is only available in transformers>=4.50.0
from transformers.models.gpt2.tokenization_gpt2_fast import GPT2TokenizerFast
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
from transformers.models.smolvlm.configuration_smolvlm import SmolVLMConfig
from transformers.models.smolvlm.image_processing_smolvlm import SmolVLMImageProcessor
from transformers.models.smolvlm.modeling_smolvlm import SmolVLMForConditionalGeneration
@@ -209,7 +228,6 @@ try:
except ImportError:
NUM2WORDS_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -268,12 +286,12 @@ if LLAMA4_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
),
attn_implementation="sdpa",
pad_token_id=None,
),
)
@@ -321,8 +339,9 @@ if MLLAMA_AVAILABLE:
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
rope_theta=500_000,
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -372,7 +391,6 @@ if PALIGEMMA_AVAILABLE:
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -421,7 +439,6 @@ if PALIGEMMA_AVAILABLE:
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -466,7 +483,6 @@ if GEMMA3_AVAILABLE:
rms_norm_eps=1e-06,
use_cache=True,
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
).to_dict(),
@@ -503,10 +519,10 @@ if QWEN2_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=True,
@@ -545,8 +561,6 @@ if LLAVA_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
pretraining_tp=1,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
max_position_embeddings=4096, # llava-1.5-7b-hf
@@ -637,7 +651,6 @@ if SMOLVLM2_AVAILABLE:
num_hidden_layers=4, # 30 -> reduced to 4 for testing
num_key_value_heads=3, # 3 for 256M model
rms_norm_eps=1e-5,
rope_theta=100000,
tie_word_embeddings=False,
vocab_size=49280,
),
@@ -680,10 +693,10 @@ if QWEN2_5_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=True,
@@ -742,11 +755,12 @@ if QWEN3_VL_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
attention_dropout=0.0,
attention_bias=False,
).to_dict(),
@@ -794,11 +808,12 @@ if QWEN3_VL_MOE_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
attention_dropout=0.0,
attention_bias=False,
decoder_sparse_step=1,
@@ -806,6 +821,7 @@ if QWEN3_VL_MOE_AVAILABLE:
num_experts_per_tok=2,
num_experts=4,
mlp_only_layers=[],
pad_token_id=None,
).to_dict(),
),
)
@@ -825,7 +841,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Qwen2VLImageProcessor()
video_processor = Qwen2VLVideoProcessor()
return Qwen2VLProcessor(
@@ -847,7 +863,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Qwen2VLImageProcessor()
video_processor = Qwen2VLVideoProcessor()
return Qwen2_5_VLProcessor(
@@ -869,7 +885,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Qwen2VLImageProcessor(patch_size=16, temporal_patch_size=2, merge_size=2)
video_processor = Qwen3VLVideoProcessor()
return Qwen3VLProcessor(
@@ -926,7 +942,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = GotOcr2ImageProcessorFast(
crop_to_patches=False, min_patches=1, max_patches=12, size={"height": 448, "width": 448}
)
@@ -950,7 +966,7 @@ def create_processor(model_name: str):
)
]
)
gpt2_tokenizer = GPT2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
gpt2_tokenizer = GPT2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = SmolVLMImageProcessor(size={"longest_edge": 512})
video_processor = SmolVLMVideoProcessor()
@@ -1020,7 +1036,7 @@ def create_processor(model_name: str):
]
)
fast_tokenizer = GemmaTokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = SiglipImageProcessor(size={"height": 224, "width": 224}, image_seq_length=256)
return PaliGemmaProcessor(image_processor=image_processor, tokenizer=fast_tokenizer)
@@ -1040,7 +1056,7 @@ def create_processor(model_name: str):
)
]
)
fast_tokenizer = GemmaTokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Gemma3ImageProcessor()
return Gemma3Processor(image_processor=image_processor, tokenizer=fast_tokenizer)
@@ -1402,6 +1418,11 @@ def run_mini_model_multimodal(
not LLAMA4_AVAILABLE,
reason="Llama4 not available in this version of transformers",
),
# TODO: Remove this skipif when the bug fix is released in Transformers
pytest.mark.skipif(
version.parse(transformers.__version__) <= version.parse("5.1.0"),
reason="Wait for this bug fix to be released in Transformers: https://github.com/huggingface/transformers/pull/43882",
),
],
),
pytest.param(
test/convergence/bf16/test_mini_models_with_logits.py
+42 -76
View File
@@ -4,8 +4,10 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # Ensure deterministic behavi
import pytest
import torch
import transformers
from datasets import load_from_disk
from packaging import version
from torch.utils.data import DataLoader
from transformers.models.gemma import GemmaConfig
from transformers.models.gemma import GemmaForCausalLM
@@ -52,6 +54,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_next
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl_moe
from liger_kernel.transformers import apply_liger_kernel_to_smollm3
from liger_kernel.utils import infer_device
from test.utils import DEFAULT_DATASET_PATH
from test.utils import MiniModelConfig
from test.utils import assert_verbose_allclose
@@ -92,6 +95,8 @@ from test.utils import set_seed
from test.utils import simple_collate_fn
from test.utils import supports_bfloat16
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
try:
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM
@@ -289,8 +294,6 @@ except ImportError:
EXAONE4_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
MINI_MODEL_SETUPS = {
@@ -315,8 +318,6 @@ MINI_MODEL_SETUPS = {
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -344,7 +345,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -373,7 +373,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=None, # defaults to num_attention_heads
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=None,
tie_word_embeddings=False,
use_cache=True,
@@ -398,7 +397,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -423,7 +421,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -458,7 +455,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -486,7 +482,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -514,7 +509,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -543,8 +537,6 @@ if LLAMA4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=10000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -571,7 +563,6 @@ if QWEN3_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -597,8 +588,6 @@ if QWEN3_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -642,15 +631,16 @@ if QWEN3_VL_AVAILABLE:
num_key_value_heads=2,
pad_token_id=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
sliding_window=131072,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
vision_config=dict(
depth=4,
@@ -697,11 +687,6 @@ if QWEN3_VL_MOE_AVAILABLE:
num_key_value_heads=2,
pad_token_id=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
sliding_window=131072,
tie_word_embeddings=False,
use_cache=True,
@@ -751,7 +736,6 @@ if GEMMA3_AVAILABLE:
bos_token_id=2,
eos_token_id=1,
tie_word_embeddings=True,
rope_theta=10000.0, # 1000000
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -778,18 +762,20 @@ if MLLAMA_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_scaling=dict(
factor=8.0,
high_freq_factor=4.0,
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_theta=500_000,
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
)
@@ -819,10 +805,10 @@ if QWEN2_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=False,
@@ -871,10 +857,10 @@ if QWEN2_5_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=False,
@@ -923,8 +909,6 @@ if GRANITE_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -957,8 +941,6 @@ if LLAVA_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
pretraining_tp=1,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
max_position_embeddings=4096, # llava-1.5-7b-hf
@@ -1016,8 +998,6 @@ if OLMO2_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1045,8 +1025,6 @@ if OLMO3_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1075,8 +1053,6 @@ if GLM4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1111,8 +1087,6 @@ if GLM4V_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1128,13 +1102,14 @@ if GLM4V_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
"pad_token_id": None,
},
vision_config={
"depth": 4, # 32
@@ -1178,8 +1153,6 @@ if GLM4V_MOE_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1195,11 +1168,6 @@ if GLM4V_MOE_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
"attention_dropout": 0.0,
@@ -1212,6 +1180,11 @@ if GLM4V_MOE_AVAILABLE:
"topk_group": 1,
"first_k_dense_replace": 1,
"norm_topk_prob": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
},
vision_config={
"depth": 4, # 32
@@ -1249,8 +1222,6 @@ if SMOLLM3_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1341,8 +1312,6 @@ if QWEN3NEXT_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -1383,7 +1352,6 @@ if HUNYUAN_V1_AVAILABLE:
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
@@ -1414,8 +1382,6 @@ if HUNYUAN_V1_AVAILABLE:
eod_token_id=3,
sep_token_id=4,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
attention_dropout=0.0,
num_experts=2,
@@ -1442,11 +1408,11 @@ if EXAONE4_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=1000000.0,
tie_word_embeddings=True,
use_cache=True,
vocab_size=32000,
attn_implementation="sdpa",
pad_token_id=None,
),
)
@@ -1540,12 +1506,12 @@ def run_mini_model(
[
# Tolerance is set higher than usual to pass the tests.
pytest.param(
"mini_llama4",
"mini_llama4", # llama4 requires slightly larger tolerances to pass this test after bug fix to llama4 in transformers v5.0.0
32,
1e-5,
torch.bfloat16,
1e-2,
5e-2,
4e-1,
3e-1,
2e-1,
1e-2,
@@ -1666,7 +1632,7 @@ def run_mini_model(
1e-5,
torch.bfloat16,
1e-2,
5e-2,
2e-1,
1e-1,
1e-2,
1e-2,
test/convergence/fp32/test_mini_models.py
+97 -133
View File
@@ -4,8 +4,10 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # Ensure deterministic behavi
import pytest
import torch
import transformers
from datasets import load_from_disk
from packaging import version
from torch.utils.data import DataLoader
from transformers.models.gemma import GemmaConfig
from transformers.models.gemma import GemmaForCausalLM
@@ -53,6 +55,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_next
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl_moe
from liger_kernel.transformers import apply_liger_kernel_to_smollm3
from liger_kernel.utils import infer_device
from test.utils import DEFAULT_DATASET_PATH
from test.utils import MiniModelConfig
from test.utils import assert_verbose_allclose
@@ -93,6 +96,8 @@ from test.utils import revert_liger_kernel_to_smollm3
from test.utils import set_seed
from test.utils import simple_collate_fn
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
try:
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM
@@ -306,8 +311,6 @@ except ImportError:
EXAONE4_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
MINI_MODEL_SETUPS = {
@@ -332,8 +335,6 @@ MINI_MODEL_SETUPS = {
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -361,7 +362,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -390,7 +390,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=None, # defaults to num_attention_heads
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=None,
tie_word_embeddings=False,
use_cache=True,
@@ -415,7 +414,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -440,7 +438,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -475,7 +472,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -503,7 +499,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -531,7 +526,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -559,8 +553,6 @@ if LLAMA4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=10000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -588,7 +580,6 @@ if QWEN3_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -614,8 +605,6 @@ if QWEN3_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -651,14 +640,6 @@ if GPT_OSS_AVAILABLE:
rms_norm_eps=1e-5,
use_cache=True,
tie_word_embeddings=False,
rope_parameters={
"rope_type": "yarn",
"factor": 8.0,
"beta_fast": 32.0,
"beta_slow": 1.0,
"truncate": False,
"original_max_position_embeddings": 4096,
},
attention_dropout=0.0,
num_local_experts=8, # Reduced from 32 for mini model
num_experts_per_tok=2, # Reduced from 4 for mini model
@@ -691,7 +672,6 @@ if GEMMA3_AVAILABLE:
bos_token_id=2,
eos_token_id=1,
tie_word_embeddings=True,
rope_theta=10000.0, # 1000000
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -718,18 +698,20 @@ if MLLAMA_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_scaling=dict(
factor=8.0,
high_freq_factor=4.0,
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_theta=500_000,
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
)
@@ -739,36 +721,38 @@ if QWEN2_VL_AVAILABLE:
liger_kernel_patch_revert_func=revert_liger_kernel_to_qwen2_vl,
model_class=Qwen2VLForConditionalGeneration,
mini_model_config=Qwen2VLConfig(
attention_dropout=0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
bos_token_id=1, # 151643
eos_token_id=2, # 151645
# In transformers v5, text-related parameters must be in text_config
text_config={
"attention_dropout": 0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
"bos_token_id": 1, # 151643
"eos_token_id": 2, # 151645
"hidden_act": "silu",
"hidden_size": 1536, # 8192
"initializer_range": 0.02,
"intermediate_size": 4864, # 29568
"max_position_embeddings": 32768,
"max_window_layers": 4, # 80
"num_attention_heads": 12, # 64
"num_hidden_layers": 4, # 80
"num_key_value_heads": 2, # 8
"rms_norm_eps": 1e-6, # 1e-5
**(
{"rope_parameters": {"mrope_section": [16, 24, 24]}} # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else {"rope_scaling": {"type": "mrope", "mrope_section": [16, 24, 24]}}
),
"sliding_window": 4096,
"tie_word_embeddings": False,
"use_cache": True,
"vocab_size": 32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
"use_sliding_window": False,
},
vision_start_token_id=32765, # vocab_size - 5
vision_end_token_id=32766, # vocab_size - 4
vision_token_id=32767, # vocab_size - 3
image_token_id=32768, # vocab_size - 2
video_token_id=32769, # vocab_size - 1
hidden_act="silu",
hidden_size=1536, # 8192
initializer_range=0.02,
intermediate_size=4864, # 29568
max_position_embeddings=32768,
max_window_layers=4, # 80
num_attention_heads=12, # 64
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
),
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
use_sliding_window=False,
vision_config={
"depth": 4, # 32
"embed_dim": 1280,
@@ -781,7 +765,6 @@ if QWEN2_VL_AVAILABLE:
"spatial_patch_size": 14,
"temporal_patch_size": 2,
},
attn_implementation="sdpa",
),
)
@@ -791,36 +774,38 @@ if QWEN2_5_VL_AVAILABLE:
liger_kernel_patch_revert_func=revert_liger_kernel_to_qwen2_5_vl,
model_class=Qwen2_5_VLForConditionalGeneration,
mini_model_config=Qwen2_5_VLConfig(
attention_dropout=0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
bos_token_id=1, # 151643
eos_token_id=2, # 151645
# In transformers v5, text-related parameters must be in text_config
text_config={
"attention_dropout": 0.0,
# bos and eos set to match the Mistral-7B tokenizer used to create the test dataset
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
"bos_token_id": 1, # 151643
"eos_token_id": 2, # 151645
"hidden_act": "silu",
"hidden_size": 1536, # 8192
"initializer_range": 0.02,
"intermediate_size": 4864, # 29568
"max_position_embeddings": 32768,
"max_window_layers": 4, # 80
"num_attention_heads": 12, # 64
"num_hidden_layers": 4, # 80
"num_key_value_heads": 2, # 8
"rms_norm_eps": 1e-6, # 1e-5
**(
{"rope_parameters": {"mrope_section": [16, 24, 24]}} # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else {"rope_scaling": {"type": "mrope", "mrope_section": [16, 24, 24]}}
),
"sliding_window": 4096,
"tie_word_embeddings": False,
"use_cache": True,
"vocab_size": 32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
"use_sliding_window": False,
},
vision_start_token_id=32765, # vocab_size - 5
vision_end_token_id=32766, # vocab_size - 4
vision_token_id=32767, # vocab_size - 3
image_token_id=32768, # vocab_size - 2
video_token_id=32769, # vocab_size - 1
hidden_act="silu",
hidden_size=1536, # 8192
initializer_range=0.02,
intermediate_size=4864, # 29568
max_position_embeddings=32768,
max_window_layers=4, # 80
num_attention_heads=12, # 64
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
),
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32768, # 152064 # >32k, Mistral-7B tokenizer vocab size
use_sliding_window=False,
vision_config={
"depth": 4, # 32
"hidden_act": "silu",
@@ -837,7 +822,6 @@ if QWEN2_5_VL_AVAILABLE:
"tokens_per_second": 2,
"temporal_patch_size": 2,
},
attn_implementation="sdpa",
),
)
@@ -866,13 +850,14 @@ if QWEN3_VL_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
use_cache=True,
vocab_size=32768,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
vision_config=dict(
depth=4,
@@ -919,11 +904,6 @@ if QWEN3_VL_MOE_AVAILABLE:
num_key_value_heads=2,
head_dim=128,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
use_cache=True,
vocab_size=32768,
decoder_sparse_step=1,
@@ -932,6 +912,13 @@ if QWEN3_VL_MOE_AVAILABLE:
num_experts=4,
tie_word_embeddings=False,
mlp_only_layers=[],
pad_token_id=None,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
).to_dict(),
vision_config=Qwen3VLMoeVisionConfig(
depth=4,
@@ -973,8 +960,6 @@ if GRANITE_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1006,8 +991,6 @@ if OLMO2_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1035,8 +1018,6 @@ if OLMO3_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1065,8 +1046,6 @@ if GLM4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1102,8 +1081,6 @@ if GLM4V_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1119,13 +1096,14 @@ if GLM4V_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
"pad_token_id": None,
},
vision_config={
"depth": 4, # 32
@@ -1169,8 +1147,6 @@ if GLM4V_MOE_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1186,11 +1162,6 @@ if GLM4V_MOE_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
"attention_dropout": 0.0,
@@ -1203,6 +1174,11 @@ if GLM4V_MOE_AVAILABLE:
"topk_group": 1,
"first_k_dense_replace": 1,
"norm_topk_prob": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
},
vision_config={
"depth": 4, # 32
@@ -1238,8 +1214,6 @@ if LLAVA_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
pretraining_tp=1,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
max_position_embeddings=4096, # llava-1.5-7b-hf
@@ -1298,8 +1272,6 @@ if SMOLLM3_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1390,8 +1362,6 @@ if QWEN3NEXT_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -1430,7 +1400,6 @@ if HUNYUAN_V1_AVAILABLE:
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
@@ -1456,7 +1425,6 @@ if HUNYUAN_V1_AVAILABLE:
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
num_experts=8,
@@ -1484,11 +1452,11 @@ if EXAONE4_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=1000000.0,
tie_word_embeddings=True,
use_cache=True,
vocab_size=32000,
attn_implementation="sdpa",
pad_token_id=None,
),
)
@@ -1559,7 +1527,6 @@ def run_mini_model(
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
loss_list = []
for i in range(num_steps):
batch = next(loader_iter).to(model.device)
optimizer.zero_grad()
@@ -1590,12 +1557,12 @@ def run_mini_model(
"model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
[
pytest.param(
"mini_llama4",
"mini_llama4", # llama4 requires slightly larger tolerances to pass this test after bug fix to llama4 in transformers v5.0.0
32,
1e-4,
torch.float32,
1e-8,
1e-5,
1e-3,
5e-3,
1e-3,
5e-3,
@@ -1876,10 +1843,7 @@ def run_mini_model(
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
version.parse(transformers.__version__) < version.parse("4.49.0"),
reason="Mistral not available in transformers<=4.49.0",
),
marks=[],
),
# 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),
test/convergence/fp32/test_mini_models_multimodal.py
+48 -52
View File
@@ -5,11 +5,12 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # Ensure deterministic behavi
import pytest
import torch
import transformers
from datasets import load_dataset
from packaging import version
from torch.utils.data import DataLoader
from transformers import PreTrainedTokenizerFast
from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
from liger_kernel.transformers import apply_liger_kernel_to_gemma3
@@ -23,6 +24,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen2_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl_moe
from liger_kernel.transformers import apply_liger_kernel_to_smolvlm
from liger_kernel.utils import infer_device
from test.utils import FAKE_CONFIGS_PATH
from test.utils import UNTOKENIZED_DATASET_PATH
from test.utils import MiniModelConfig
@@ -49,12 +51,23 @@ from test.utils import revert_liger_kernel_to_smolvlm2
from test.utils import set_seed
from test.utils import train_bpe_tokenizer
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.gemma.tokenization_gemma import GemmaTokenizer
else:
from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast as GemmaTokenizer
try:
# Qwen2-VL is only available in transformers>=4.52.4
import transformers
from packaging import version
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig
from transformers.models.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLForConditionalGeneration
@@ -70,7 +83,11 @@ try:
import transformers
from packaging import version
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
@@ -83,7 +100,10 @@ except ImportError:
try:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
@@ -108,7 +128,10 @@ except ImportError:
QWEN3_VL_MOE_AVAILABLE = False
try:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
else:
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast as Qwen2Tokenizer
from transformers.models.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
@@ -183,7 +206,7 @@ try:
from transformers.models.paligemma.processing_paligemma import PaliGemmaProcessor
from transformers.models.siglip.image_processing_siglip import SiglipImageProcessor
PALIGEMMA_AVAILABLE = version.parse(transformers.__version__) >= version.parse("4.46.0")
PALIGEMMA_AVAILABLE = True
except ImportError:
PALIGEMMA_AVAILABLE = False
@@ -214,7 +237,7 @@ except ImportError:
try:
# SmolVLM2 is only available in transformers>=4.50.0
from transformers.models.gpt2.tokenization_gpt2_fast import GPT2TokenizerFast
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
from transformers.models.smolvlm.configuration_smolvlm import SmolVLMConfig
from transformers.models.smolvlm.image_processing_smolvlm import SmolVLMImageProcessor
from transformers.models.smolvlm.modeling_smolvlm import SmolVLMForConditionalGeneration
@@ -232,7 +255,6 @@ try:
except ImportError:
NUM2WORDS_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -291,7 +313,6 @@ if LLAMA4_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -345,8 +366,9 @@ if MLLAMA_AVAILABLE:
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
rope_theta=500_000,
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -396,7 +418,6 @@ if PALIGEMMA_AVAILABLE:
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -446,7 +467,6 @@ if PALIGEMMA_AVAILABLE:
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -492,7 +512,6 @@ if GEMMA3_AVAILABLE:
rms_norm_eps=1e-06,
use_cache=True,
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -528,10 +547,10 @@ if QWEN2_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=True,
@@ -570,8 +589,6 @@ if LLAVA_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
pretraining_tp=1,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
max_position_embeddings=4096, # llava-1.5-7b-hf
@@ -662,7 +679,6 @@ if SMOLVLM2_AVAILABLE:
num_hidden_layers=4, # 30 -> reduced to 4 for testing
num_key_value_heads=3, # 3 for 256M model
rms_norm_eps=1e-5,
rope_theta=100000,
tie_word_embeddings=False,
vocab_size=49280,
),
@@ -705,10 +721,10 @@ if QWEN2_5_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=True,
@@ -767,11 +783,6 @@ if QWEN3_VL_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
).to_dict(),
@@ -819,11 +830,6 @@ if QWEN3_VL_MOE_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
decoder_sparse_step=1,
@@ -879,11 +885,6 @@ if QWEN3_VL_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
).to_dict(),
@@ -931,11 +932,6 @@ if QWEN3_VL_MOE_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
decoder_sparse_step=1,
@@ -962,7 +958,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Qwen2VLImageProcessor()
video_processor = Qwen2VLVideoProcessor()
return Qwen2VLProcessor(
@@ -984,7 +980,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Qwen2VLImageProcessor()
video_processor = Qwen2VLVideoProcessor()
return Qwen2_5_VLProcessor(
@@ -1006,7 +1002,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Qwen2VLImageProcessor(patch_size=16, temporal_patch_size=2, merge_size=2)
video_processor = Qwen3VLVideoProcessor()
return Qwen3VLProcessor(
@@ -1063,7 +1059,7 @@ def create_processor(model_name: str):
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
qwen_tokenizer = Qwen2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = GotOcr2ImageProcessorFast(
crop_to_patches=False, min_patches=1, max_patches=12, size={"height": 448, "width": 448}
)
@@ -1087,7 +1083,7 @@ def create_processor(model_name: str):
)
]
)
gpt2_tokenizer = GPT2TokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
gpt2_tokenizer = GPT2Tokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = SmolVLMImageProcessor(size={"longest_edge": 512})
video_processor = SmolVLMVideoProcessor()
@@ -1157,7 +1153,7 @@ def create_processor(model_name: str):
)
]
)
fast_tokenizer = GemmaTokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = SiglipImageProcessor(size={"height": 224, "width": 224}, image_seq_length=256)
return PaliGemmaProcessor(image_processor=image_processor, tokenizer=fast_tokenizer)
@@ -1177,7 +1173,7 @@ def create_processor(model_name: str):
)
]
)
fast_tokenizer = GemmaTokenizerFast(tokenizer_object=tokenizer_base, **tokenizer_config)
fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
image_processor = Gemma3ImageProcessor()
return Gemma3Processor(image_processor=image_processor, tokenizer=fast_tokenizer)
test/convergence/fp32/test_mini_models_with_logits.py
+41 -74
View File
@@ -4,8 +4,10 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # Ensure deterministic behavi
import pytest
import torch
import transformers
from datasets import load_from_disk
from packaging import version
from torch.utils.data import DataLoader
from transformers.models.gemma import GemmaConfig
from transformers.models.gemma import GemmaForCausalLM
@@ -52,6 +54,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_next
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_vl_moe
from liger_kernel.transformers import apply_liger_kernel_to_smollm3
from liger_kernel.utils import infer_device
from test.utils import DEFAULT_DATASET_PATH
from test.utils import MiniModelConfig
from test.utils import assert_verbose_allclose
@@ -91,6 +94,8 @@ from test.utils import revert_liger_kernel_to_smollm3
from test.utils import set_seed
from test.utils import simple_collate_fn
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
try:
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM
@@ -309,8 +314,6 @@ except ImportError:
EXAONE4_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
MINI_MODEL_SETUPS = {
@@ -335,8 +338,6 @@ MINI_MODEL_SETUPS = {
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -364,7 +365,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -393,7 +393,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=None, # defaults to num_attention_heads
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=None,
tie_word_embeddings=False,
use_cache=True,
@@ -418,7 +417,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -443,7 +441,6 @@ MINI_MODEL_SETUPS = {
num_hidden_layers=4, # 32
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_theta=10000.0,
sliding_window=4096,
tie_word_embeddings=False,
use_cache=True,
@@ -478,7 +475,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -506,7 +502,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
),
@@ -534,7 +529,6 @@ MINI_MODEL_SETUPS = {
bos_token_id=1, # 128000
eos_token_id=2, # 128001
tie_word_embeddings=True,
rope_theta=10000.0,
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -562,8 +556,6 @@ if LLAMA4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=10000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -590,7 +582,6 @@ if QWEN3_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
sliding_window=131072,
tie_word_embeddings=True,
use_cache=True,
@@ -616,8 +607,6 @@ if QWEN3_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -656,7 +645,6 @@ if GEMMA3_AVAILABLE:
bos_token_id=2,
eos_token_id=1,
tie_word_embeddings=True,
rope_theta=10000.0, # 1000000
attention_bias=False,
attention_dropout=0.0,
attn_implementation="eager",
@@ -683,18 +671,20 @@ if MLLAMA_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_scaling=dict(
factor=8.0,
high_freq_factor=4.0,
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
rope_theta=500_000,
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
)
@@ -724,10 +714,10 @@ if QWEN2_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=False,
@@ -776,10 +766,10 @@ if QWEN2_5_VL_AVAILABLE:
num_hidden_layers=4, # 80
num_key_value_heads=2, # 8
rms_norm_eps=1e-6, # 1e-5
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
**(
dict(rope_parameters=dict(mrope_section=[16, 24, 24])) # (temporal, height, width)
if IS_TRANSFORMERS_V5_OR_LATER
else dict(rope_scaling=dict(type="mrope", mrope_section=[16, 24, 24]))
),
sliding_window=4096,
tie_word_embeddings=False,
@@ -833,15 +823,16 @@ if QWEN3_VL_AVAILABLE:
num_key_value_heads=2,
pad_token_id=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
sliding_window=131072,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24], # (temporal, height, width)
)
if not IS_TRANSFORMERS_V5_OR_LATER
else None,
),
vision_config=dict(
depth=4,
@@ -888,11 +879,6 @@ if QWEN3_VL_MOE_AVAILABLE:
num_key_value_heads=2,
pad_token_id=2,
rms_norm_eps=1e-6,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
sliding_window=131072,
tie_word_embeddings=False,
use_cache=True,
@@ -942,8 +928,6 @@ if GRANITE_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -976,8 +960,6 @@ if LLAVA_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
pretraining_tp=1,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
max_position_embeddings=4096, # llava-1.5-7b-hf
@@ -1035,8 +1017,6 @@ if OLMO2_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1064,8 +1044,6 @@ if OLMO3_AVAILABLE:
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1094,8 +1072,6 @@ if GLM4_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1131,8 +1107,6 @@ if GLM4V_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1148,13 +1122,14 @@ if GLM4V_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
"pad_token_id": None,
},
vision_config={
"depth": 4, # 32
@@ -1197,8 +1172,6 @@ if GLM4V_MOE_AVAILABLE:
num_hidden_layers=4, # 61
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 151552
@@ -1214,11 +1187,6 @@ if GLM4V_MOE_AVAILABLE:
"num_hidden_layers": 4,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-5,
"rope_scaling": {
"type": "default",
"mrope_section": [8, 12, 12], # (temporal, height, width)
},
"rope_theta": 500_000,
"vocab_size": 32000,
"attention_bias": True,
"attention_dropout": 0.0,
@@ -1231,6 +1199,11 @@ if GLM4V_MOE_AVAILABLE:
"topk_group": 1,
"first_k_dense_replace": 1,
"norm_topk_prob": True,
**(
{"rope_scaling": {"type": "default", "mrope_section": [8, 12, 12]}}
if not IS_TRANSFORMERS_V5_OR_LATER
else {}
),
},
vision_config={
"depth": 4, # 32
@@ -1268,8 +1241,6 @@ if SMOLLM3_AVAILABLE:
num_key_value_heads=2, # 8
pretraining_tp=1,
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500000.0,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
@@ -1360,8 +1331,6 @@ if QWEN3NEXT_AVAILABLE:
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
use_sliding_window=False,
sliding_window=4096,
@@ -1402,7 +1371,6 @@ if HUNYUAN_V1_AVAILABLE:
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
@@ -1428,7 +1396,6 @@ if HUNYUAN_V1_AVAILABLE:
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
num_experts=8,
@@ -1456,11 +1423,11 @@ if EXAONE4_AVAILABLE:
num_hidden_layers=4,
num_key_value_heads=2,
rms_norm_eps=1e-5,
rope_theta=1000000.0,
tie_word_embeddings=True,
use_cache=True,
vocab_size=32000,
attn_implementation="sdpa",
pad_token_id=None,
),
)
@@ -1553,12 +1520,12 @@ def run_mini_model(
"model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
[
pytest.param(
"mini_llama4",
"mini_llama4", # llama4 requires slightly larger tolerances to pass this test after bug fix to llama4 in transformers v5.0.0
32,
1e-4,
torch.float32,
1e-8,
1e-5,
1e-3,
5e-3,
1e-5,
5e-3,
test/transformers/test_monkey_patch.py
+92 -97
View File
@@ -10,11 +10,14 @@ import torch
import transformers
from packaging import version
from test.utils import get_mllama_rope_config
from test.utils import get_qwen3_vl_rope_config
from transformers import AutoModelForCausalLM
from transformers import PretrainedConfig
from transformers import PreTrainedModel
from liger_kernel.transformers import LigerBlockSparseTop2MLP
from liger_kernel.transformers import LigerExperts
from liger_kernel.transformers import LigerGEGLUMLP
from liger_kernel.transformers import LigerPhi3SwiGLUMLP
from liger_kernel.transformers import LigerQwen3MoeSwiGLUMLP
@@ -22,43 +25,32 @@ from liger_kernel.transformers import LigerRMSNorm
from liger_kernel.transformers import LigerSwiGLUMLP
from liger_kernel.transformers import monkey_patch
from liger_kernel.transformers.layer_norm import LigerLayerNorm
from liger_kernel.transformers.model.falcon_h1 import lce_forward as falcon_h1_lce_forward
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.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
from liger_kernel.transformers.model.paligemma import lce_forward as paligemma_lce_forward
from liger_kernel.transformers.model.phi3 import lce_forward as phi3_lce_forward
from liger_kernel.transformers.model.qwen2 import lce_forward as qwen2_lce_forward
from liger_kernel.transformers.model.qwen3_next import lce_forward as qwen3_next_lce_forward
from liger_kernel.transformers.model.smollm3 import lce_forward as smolllm3_lce_forward
from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN
from liger_kernel.transformers.monkey_patch import _apply_liger_kernel
from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance
# Import transformer version check
# We only support transformers >= 4.52.0
transformer_version = version.parse(transformers.__version__)
SUPPORTED_TRANSFORMER_VERSION = "4.46.1"
MIN_SUPPORTED_TRANSFORMERS_VERSION = version.parse("4.52.0")
if transformer_version < MIN_SUPPORTED_TRANSFORMERS_VERSION:
pytest.skip(
f"tests require transformers >= {MIN_SUPPORTED_TRANSFORMERS_VERSION}, got {transformers.__version__}",
allow_module_level=True,
)
# Import forward functions based on transformer version
if transformer_version >= version.parse(SUPPORTED_TRANSFORMER_VERSION):
from liger_kernel.transformers.model.falcon_h1 import lce_forward as falcon_h1_lce_forward
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.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
from liger_kernel.transformers.model.paligemma import lce_forward as paligemma_lce_forward
from liger_kernel.transformers.model.phi3 import lce_forward as phi3_lce_forward
from liger_kernel.transformers.model.qwen2 import lce_forward as qwen2_lce_forward
from liger_kernel.transformers.model.qwen3_next import lce_forward as qwen3_next_lce_forward
from liger_kernel.transformers.model.smollm3 import lce_forward as smolllm3_lce_forward
else:
from liger_kernel.transformers.model.gemma import lce_forward_deprecated as gemma_lce_forward
from liger_kernel.transformers.model.gemma2 import lce_forward_deprecated as gemma2_lce_forward
from liger_kernel.transformers.model.llama import lce_forward_deprecated as llama_lce_forward
from liger_kernel.transformers.model.mistral import (
lce_forward as mistral_lce_forward, # mistral doesn't have deprecated version
)
from liger_kernel.transformers.model.mixtral import lce_forward_deprecated as mixtral_lce_forward
from liger_kernel.transformers.model.mllama import lce_forward_deprecated as mllama_lce_forward
from liger_kernel.transformers.model.paligemma import lce_forward_deprecated as paligemma_lce_forward
from liger_kernel.transformers.model.phi3 import lce_forward_deprecated as phi3_lce_forward
from liger_kernel.transformers.model.qwen2 import lce_forward_deprecated as qwen2_lce_forward
from liger_kernel.transformers.model.qwen3_next import (
lce_forward as qwen3_next_lce_forward, # qwen3_next doesn't have deprecated version
)
IS_TRANSFORMERS_V5_OR_LATER = transformer_version >= version.parse("5.0.0")
# Check if optional modules are available
@@ -497,13 +489,9 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_for_conditional_generation(
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
**get_qwen3_vl_rope_config(), # Version-aware rope configuration
).to_dict(),
)
dummy_model_instance = Qwen3VLForConditionalGeneration._from_config(config)
@@ -598,13 +586,9 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl():
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
**get_qwen3_vl_rope_config(), # Version-aware rope configuration
).to_dict(),
)
dummy_model_instance = Qwen3VLModel._from_config(config)
@@ -675,13 +659,9 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_text():
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
**get_qwen3_vl_rope_config(), # Version-aware rope configuration
)
dummy_model_instance = Qwen3VLTextModel._from_config(config)
@@ -771,11 +751,6 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_moe_for_conditional_generat
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
decoder_sparse_step=1,
@@ -783,6 +758,8 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_moe_for_conditional_generat
num_experts_per_tok=2,
num_experts=4,
mlp_only_layers=[],
pad_token_id=None,
**get_qwen3_vl_rope_config(), # Version-aware rope configuration
).to_dict(),
)
dummy_model_instance = Qwen3VLMoeForConditionalGeneration._from_config(config)
@@ -877,11 +854,6 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_moe():
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
decoder_sparse_step=1,
@@ -889,6 +861,8 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_moe():
num_experts_per_tok=2,
num_experts=4,
mlp_only_layers=[],
pad_token_id=None,
**get_qwen3_vl_rope_config(), # Version-aware rope configuration
).to_dict(),
)
dummy_model_instance = Qwen3VLMoeModel._from_config(config)
@@ -959,11 +933,6 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_moe_text():
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=True,
rope_theta=1000000.0,
rope_scaling=dict(
type="mrope",
mrope_section=[16, 24, 24],
),
attention_dropout=0.0,
attention_bias=False,
decoder_sparse_step=1,
@@ -971,6 +940,8 @@ def test_apply_liger_kernel_to_instance_for_qwen3_vl_moe_text():
num_experts_per_tok=2,
num_experts=4,
mlp_only_layers=[],
pad_token_id=None,
**get_qwen3_vl_rope_config(), # Version-aware rope configuration
)
dummy_model_instance = Qwen3VLMoeTextModel._from_config(config)
@@ -1106,13 +1077,7 @@ def test_apply_liger_kernel_to_instance_for_mllama_for_conditional_generation():
intermediate_size=64,
hidden_act="silu",
num_hidden_layers=2,
rope_scaling=dict(
factor=8.0,
high_freq_factor=4.0,
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
**get_mllama_rope_config(), # Version-aware rope configuration
),
vision_config=transformers.models.mllama.configuration_mllama.MllamaVisionConfig(
rms_norm_eps=1e-5,
@@ -1204,13 +1169,7 @@ def test_apply_liger_kernel_to_instance_for_mllama_for_causal_lm():
intermediate_size=64,
hidden_act="silu",
num_hidden_layers=2,
rope_scaling=dict(
factor=8.0,
high_freq_factor=4.0,
low_freq_factor=1.0,
original_max_position_embeddings=8192,
rope_type="llama3",
),
**get_mllama_rope_config(), # Version-aware rope configuration
)
dummy_model_instance = MllamaForCausalLM._from_config(config)
@@ -1319,6 +1278,7 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_conditional_generation():
num_hidden_layers=2,
vision_output_dim=64,
),
pad_token_id=None,
)
dummy_model_instance = Llama4ForConditionalGeneration._from_config(config)
@@ -1386,10 +1346,6 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_conditional_generation():
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
@pytest.mark.skipif(
transformer_version < version.parse("4.49.0"),
reason="fused linear cross entropy patch doesn't work on mistral in transformers<4.49.0",
)
def test_apply_liger_kernel_to_instance_for_mistral():
# Ensure any monkey patching is cleaned up for subsequent tests
with patch("transformers.models.mistral.modeling_mistral"):
@@ -1449,8 +1405,11 @@ def test_apply_liger_kernel_to_instance_for_mixtral():
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(mixtral_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
for expert in layer.block_sparse_moe.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerBlockSparseTop2MLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward)
else:
for expert in layer.block_sparse_moe.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerBlockSparseTop2MLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
@@ -1461,8 +1420,11 @@ def test_apply_liger_kernel_to_instance_for_mixtral():
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(mixtral_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
for expert in layer.block_sparse_moe.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerBlockSparseTop2MLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward)
else:
for expert in layer.block_sparse_moe.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerBlockSparseTop2MLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
@@ -1862,7 +1824,11 @@ def test_apply_liger_kernel_to_instance_for_qwen3_moe():
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_moe_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(LigerQwen3MoeSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward)
else:
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
@@ -1873,8 +1839,11 @@ def test_apply_liger_kernel_to_instance_for_qwen3_moe():
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_moe_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
for mlp_expert in layer.mlp.experts:
assert inspect.getsource(mlp_expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward)
else:
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
@@ -2596,6 +2565,7 @@ def test_apply_liger_kernel_to_instance_for_glm4v():
"hidden_size": 32,
"intermediate_size": 64,
"hidden_act": "silu",
"pad_token_id": None,
},
vision_config={
"num_hidden_layers": 2,
@@ -2703,8 +2673,11 @@ def test_apply_liger_kernel_to_instance_for_glm4v_moe():
LigerRMSNormForGlm4.forward
)
if decoder_layer.mlp.experts is not None:
for expert in decoder_layer.mlp.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(decoder_layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward)
else:
for expert in decoder_layer.mlp.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerSwiGLUMLP.forward)
if decoder_layer.mlp.shared_experts is not None:
assert inspect.getsource(decoder_layer.mlp.shared_experts.forward) != inspect.getsource(
LigerSwiGLUMLP.forward
@@ -2739,8 +2712,13 @@ def test_apply_liger_kernel_to_instance_for_glm4v_moe():
LigerRMSNormForGlm4.forward
)
if getattr(decoder_layer.mlp, "experts", None) is not None:
for expert in decoder_layer.mlp.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(decoder_layer.mlp.experts.forward) == inspect.getsource(
LigerExperts.forward
)
else:
for expert in decoder_layer.mlp.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerSwiGLUMLP.forward)
if getattr(decoder_layer.mlp, "shared_experts", None) is not None:
assert inspect.getsource(decoder_layer.mlp.shared_experts.forward) == inspect.getsource(
LigerSwiGLUMLP.forward
@@ -2822,8 +2800,15 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next():
assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward)
else:
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
if hasattr(layer.mlp, "shared_expert"):
assert inspect.getsource(layer.mlp.shared_expert.forward) != inspect.getsource(
LigerSwiGLUMLP.forward
)
else:
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
@@ -2838,8 +2823,11 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next():
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward)
else:
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
if hasattr(layer.mlp, "shared_expert"):
assert inspect.getsource(layer.mlp.shared_expert.forward) == inspect.getsource(
LigerSwiGLUMLP.forward
@@ -2878,7 +2866,11 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_moe():
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(hunyuan_v1_moe_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(LigerQwen3MoeSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward)
else:
for expert in layer.mlp.experts:
assert inspect.getsource(expert.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
@@ -2889,8 +2881,11 @@ def test_apply_liger_kernel_to_instance_for_hunyuan_v1_moe():
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(hunyuan_v1_moe_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
for mlp_expert in layer.mlp.experts:
assert inspect.getsource(mlp_expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
if IS_TRANSFORMERS_V5_OR_LATER:
assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward)
else:
for mlp_expert in layer.mlp.experts:
assert inspect.getsource(mlp_expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.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/transformers/test_rope.py
+1 -1
View File
@@ -83,7 +83,7 @@ def test_correctness(
cos, sin = rotary_emb(k1, pos_ids)
# validate forward pass
hf_q, hf_k = apply_rotary_pos_emb(q1, k1, cos, sin, pos_ids)
hf_q, hf_k = apply_rotary_pos_emb(q1, k1, cos, sin)
tt_q, tt_k = liger_rotary_pos_emb(q2, k2, cos, sin)
assert torch.allclose(hf_q, tt_q, atol=atol, rtol=rtol)
assert torch.allclose(hf_k, tt_k, atol=atol, rtol=rtol)
test/transformers/test_swiglu.py
+193
View File
@@ -1,18 +1,29 @@
import pytest
import torch
import transformers
from packaging import version
from test.utils import supports_bfloat16
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaMLP
from transformers.models.mixtral.configuration_mixtral import MixtralConfig
from transformers.models.phi3.configuration_phi3 import Phi3Config
from transformers.models.phi3.modeling_phi3 import Phi3MLP
from liger_kernel.ops.swiglu import LigerSiLUMulFunction
from liger_kernel.transformers.functional import liger_swiglu
from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP
from liger_kernel.transformers.swiglu import LigerExperts
from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
from liger_kernel.utils import infer_device
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
if IS_TRANSFORMERS_V5_OR_LATER:
from transformers.models.mixtral.modeling_mixtral import MixtralExperts
else:
from transformers.models.mixtral.modeling_mixtral import MixtralBlockSparseTop2MLP
device = infer_device()
LLAMA_CONFIG = LlamaConfig(
@@ -104,6 +115,188 @@ def test_correctness_llamamlp(bsz, seq_len, hidden_size, intermediate_size, dtyp
assert torch.allclose(x1.grad, x2.grad, atol=atol, rtol=rtol)
@pytest.mark.skipif(IS_TRANSFORMERS_V5_OR_LATER, reason="Skip for transformers >= v5.0.0")
@pytest.mark.parametrize(
"bsz, seq_len, hidden_size, intermediate_size",
[
(2, 256, 256, 512),
# weird shapes
(6, 42, 123, 431),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
# atol is for small values: they have more difference, so set atol higher
# rtol is for larger values: they are very close, so set rtol lower
(torch.float32, 1e-0, 1e-5),
# TODO: we should find a better way to tune this. 1e4 is too large apparently
pytest.param(
torch.bfloat16,
1e4,
1e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
],
)
def test_correctness_mixtralblocksparsetop2mlp(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol):
MIXTRAL_CONFIG = MixtralConfig(
num_local_experts=8,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act="silu",
num_experts_per_tok=2,
)
_input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype)
x1 = _input.clone().requires_grad_(True)
x2 = _input.clone().requires_grad_(True)
# initialize weights
G = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype)
U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype)
D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype)
mixtral_blocksparsetop2mlp = MixtralBlockSparseTop2MLP(config=MIXTRAL_CONFIG).to(device).to(dtype)
mixtral_blocksparsetop2mlp.w1.weight.data = G.T
mixtral_blocksparsetop2mlp.w2.weight.data = U.T
mixtral_blocksparsetop2mlp.w3.weight.data = D.T
liger_blocksparsetop2mlp = LigerBlockSparseTop2MLP(config=MIXTRAL_CONFIG).to(device).to(dtype)
liger_blocksparsetop2mlp.w1.weight.data = G.T
liger_blocksparsetop2mlp.w2.weight.data = U.T
liger_blocksparsetop2mlp.w3.weight.data = D.T
y1 = mixtral_blocksparsetop2mlp(x1)
y2 = liger_blocksparsetop2mlp(x2)
assert torch.allclose(y1, y2, atol=atol, rtol=rtol)
dy = torch.randn_like(y1)
y1.backward(dy.clone(), retain_graph=True)
y2.backward(dy.clone(), retain_graph=True)
assert torch.allclose(
mixtral_blocksparsetop2mlp.w1.weight.grad,
liger_blocksparsetop2mlp.w1.weight.grad,
atol=atol,
rtol=rtol,
)
assert torch.allclose(
mixtral_blocksparsetop2mlp.w2.weight.grad,
liger_blocksparsetop2mlp.w2.weight.grad,
atol=atol,
rtol=rtol,
)
assert torch.allclose(
mixtral_blocksparsetop2mlp.w3.weight.grad,
liger_blocksparsetop2mlp.w3.weight.grad,
atol=atol,
rtol=rtol,
)
assert torch.allclose(x1.grad, x2.grad, atol=atol, rtol=rtol)
@pytest.mark.skipif(not IS_TRANSFORMERS_V5_OR_LATER, reason="Skip for transformers < v5.0.0")
@pytest.mark.parametrize(
"bsz, seq_len, hidden_size, intermediate_size",
[
(2, 256, 256, 512),
# weird shapes
(6, 42, 123, 431),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
# atol is for small values: they have more difference, so set atol higher
# rtol is for larger values: they are very close, so set rtol lower
(torch.float32, 1e-0, 1e-5),
# TODO: we should find a better way to tune this. 1e4 is too large apparently
pytest.param(
torch.bfloat16,
1e4,
1e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
],
)
def test_correctness_mixtralexperts(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol):
MIXTRAL_CONFIG = MixtralConfig(
num_local_experts=8,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
experts_implementation="eager",
hidden_act="silu",
num_experts_per_tok=2,
)
_input = torch.randn(bsz * seq_len, hidden_size, device=device, dtype=dtype)
x1 = _input.clone().requires_grad_(True)
x2 = _input.clone().requires_grad_(True)
# match shape: (num_experts, 2 * intermediate_dim, hidden_dim)
GU = torch.randn(
MIXTRAL_CONFIG.num_local_experts,
2 * intermediate_size,
hidden_size,
device=device,
dtype=dtype,
requires_grad=True,
)
# match shape: (num_experts, hidden_dim, intermediate_dim)
D = torch.randn(
MIXTRAL_CONFIG.num_local_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True
)
# Generate random router logits and do topk
router_logits = torch.randn(bsz * seq_len, MIXTRAL_CONFIG.num_local_experts, device=device, dtype=dtype)
router_logits = router_logits.softmax(dim=-1)
top_k_weights, top_k_index = router_logits.topk(k=MIXTRAL_CONFIG.num_experts_per_tok, dim=-1)
top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-9)
mixtral_experts = MixtralExperts(config=MIXTRAL_CONFIG).to(device).to(dtype)
mixtral_experts.gate_up_proj.data = GU.clone().detach()
mixtral_experts.down_proj.data = D.clone().detach()
liger_experts = LigerExperts(config=MIXTRAL_CONFIG).to(device).to(dtype)
liger_experts.gate_up_proj.data = GU.clone().detach()
liger_experts.down_proj.data = D.clone().detach()
mixtral_experts.gate_up_proj.requires_grad_()
mixtral_experts.down_proj.requires_grad_()
liger_experts.gate_up_proj.requires_grad_()
liger_experts.down_proj.requires_grad_()
y1 = mixtral_experts(x1, top_k_index, top_k_weights)
y2 = liger_experts(x2, top_k_index, top_k_weights)
assert torch.allclose(y1, y2, atol=atol, rtol=rtol)
dy = torch.randn_like(y1)
y1.backward(dy.clone(), retain_graph=True)
y2.backward(dy.clone(), retain_graph=True)
assert torch.allclose(
mixtral_experts.gate_up_proj.grad,
liger_experts.gate_up_proj.grad,
atol=atol,
rtol=rtol,
)
assert torch.allclose(
mixtral_experts.down_proj.grad,
liger_experts.down_proj.grad,
atol=atol,
rtol=rtol,
)
assert torch.allclose(x1.grad, x2.grad, atol=atol, rtol=rtol)
@pytest.mark.parametrize(
"bsz, seq_len, hidden_size, intermediate_size",
[
test/transformers/test_transformers.py
+1
View File
@@ -5,6 +5,7 @@ def test_import_from_root():
try:
from liger_kernel.transformers import LigerBlockSparseTop2MLP # noqa: F401
from liger_kernel.transformers import LigerCrossEntropyLoss # noqa: F401
from liger_kernel.transformers import LigerExperts # noqa: F401
from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss # noqa: F401
from liger_kernel.transformers import LigerGEGLUMLP # noqa: F401
from liger_kernel.transformers import LigerLayerNorm # noqa: F401
test/utils.py
+61 -2
View File
@@ -31,6 +31,64 @@ from liger_kernel.utils import infer_device
device = infer_device()
# =============================================================================
# Transformers Version Compatibility Utilities
# =============================================================================
# These utilities help maintain backward compatibility across different
# versions of the transformers library (v4.52.0, v4.57.6, v5.0.0+).
TRANSFORMERS_VERSION = version.parse(transformers.__version__)
TRANSFORMERS_V5 = version.parse("5.0.0")
def is_transformers_v5_or_later() -> bool:
"""Check if the installed transformers version is 5.0.0 or later."""
return TRANSFORMERS_VERSION >= TRANSFORMERS_V5
def get_mllama_rope_config() -> dict:
"""
Get the correct rope configuration for MLlama models.
In transformers v4.x: requires explicit rope_scaling with llama3 rope_type
In transformers v5.0+: uses defaults, no explicit config needed
Returns:
dict: Configuration dictionary with rope_scaling for v4.x, empty for v5.0+
"""
if is_transformers_v5_or_later():
return {}
return {
"rope_scaling": {
"factor": 8.0,
"high_freq_factor": 4.0,
"low_freq_factor": 1.0,
"original_max_position_embeddings": 8192,
"rope_type": "llama3",
},
}
def get_qwen3_vl_rope_config() -> dict:
"""
Get the correct rope configuration for Qwen3-VL models.
In transformers v4.x: requires rope_scaling with type="mrope"
In transformers v5.0+: uses defaults, no explicit config needed
Returns:
dict: Configuration dictionary with rope_scaling for v4.x, empty for v5.0+
"""
if is_transformers_v5_or_later():
return {}
return {
"rope_theta": 1000000.0,
"rope_scaling": {
"type": "mrope",
"mrope_section": [16, 24, 24],
},
}
def set_seed(seed=42):
"""
@@ -620,11 +678,12 @@ def revert_liger_kernel_to_llava(model_config: MiniModelConfig):
Revert all Liger kernel patches applied to llava.
"""
from transformers.models.clip import modeling_clip
from transformers.models.llama import modeling_llama
from transformers.models.llava import modeling_llava
importlib.reload(modeling_clip)
# Note: Do NOT reload modeling_clip as it breaks CLIPVisionModel's
# output_hidden_states functionality in transformers v5.
# Liger kernel does not patch modeling_clip when model=None.
importlib.reload(modeling_llava)
importlib.reload(modeling_llama)