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

Monkeypatch for Llama 3.2-Vision (#282)

Browse Source 
## Summary
Add monkeypatch to support [Llama
3.2-Vision](https://github.com/huggingface/transformers/pull/33703/files#diff-d804e851851cdebeb8048938f1f8beec1cfa78bf7b1f06af86faa450f9d18def)
models.

## Details
Llama 3.2-Vision is a multimodal model. It is also only available in
transformers>=4.45.0.

Torchvision is required to run the multimodal tests for Llama 3.2-Vision
(the image processor requires it).


## Testing Done

- Hardware Type: RTX 4090
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Shao Tang <[email protected]>
Co-authored-by: Steven Shimizu <[email protected]>
Co-authored-by: shivam15s <[email protected]>
This commit is contained in:
Tyler Romero
2024-10-10 15:16:00 -07:00
committed by GitHub
co-authored by Shao Tang Steven Shimizu shivam15s
parent de12602d85
commit 9b10f48f17
13 changed files with 763 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -228,6 +228,7 @@ loss.backward()
| **Model** | **API** | **Supported Operations** |
|-------------|--------------------------------------------------------------|-------------------------------------------------------------------------|
| LLaMA 2 & 3 | `liger_kernel.transformers.apply_liger_kernel_to_llama` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| LLaMA 3.2-Vision | `liger_kernel.transformers.apply_liger_kernel_to_mllama` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| Mistral | `liger_kernel.transformers.apply_liger_kernel_to_mistral` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| Mixtral | `liger_kernel.transformers.apply_liger_kernel_to_mixtral` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| Gemma1 | `liger_kernel.transformers.apply_liger_kernel_to_gemma` | RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
pyproject.toml
+2 -1
View File
@@ -27,6 +27,7 @@ dev = [
"isort>=5.13.2",
"pytest>=7.1.2",
"datasets>=2.19.2",
"torchvision>=0.16.2",
"seaborn",
]
@@ -36,7 +37,7 @@ include = ["liger_kernel", "liger_kernel.*"]
[tool.pytest.ini_options]
pythonpath = [
"src",
"src",
"."
]
asyncio_mode = "auto"
src/liger_kernel/ops/experimental/mm_int8int2.py
-1
View File
@@ -1,5 +1,4 @@
import torch
import triton
import triton.language as tl
src/liger_kernel/transformers/__init__.py
+1
View File
@@ -16,6 +16,7 @@ from liger_kernel.transformers.monkey_patch import ( # noqa: F401
apply_liger_kernel_to_llama,
apply_liger_kernel_to_mistral,
apply_liger_kernel_to_mixtral,
apply_liger_kernel_to_mllama,
apply_liger_kernel_to_phi3,
apply_liger_kernel_to_qwen2,
apply_liger_kernel_to_qwen2_vl,
src/liger_kernel/transformers/model/mllama.py
+142
View File
@@ -0,0 +1,142 @@
from typing import List, Optional, Tuple, Union
import torch
from torch.nn import CrossEntropyLoss
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.models.mllama.modeling_mllama import MLLAMA_INPUTS_DOCSTRING
from transformers.utils import (
add_start_docstrings_to_model_forward,
replace_return_docstrings,
)
from liger_kernel.transformers.fused_linear_cross_entropy import (
LigerFusedLinearCrossEntropyLoss,
)
@add_start_docstrings_to_model_forward(MLLAMA_INPUTS_DOCSTRING)
@replace_return_docstrings(
output_type=CausalLMOutputWithPast, config_class="MllamaTextConfig"
)
def lce_forward(
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,
)
src/liger_kernel/transformers/monkey_patch.py
+81 -1
View File
@@ -117,6 +117,84 @@ def apply_liger_kernel_to_llama(
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
def apply_liger_kernel_to_mllama(
rope: bool = True,
cross_entropy: bool = False,
fused_linear_cross_entropy: bool = True,
rms_norm: bool = True,
swiglu: bool = True,
model: PreTrainedModel = None,
) -> None:
"""
Apply Liger kernels to replace original implementation in HuggingFace MLlama models.
NOTE: MLlama is not available in transformers<4.45.0
Args:
rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
fused_linear_cross_entropy (bool):
Whether to apply Liger's fused linear cross entropy loss. Default is True.
`cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
loaded. Default is None.
"""
assert not (
cross_entropy and fused_linear_cross_entropy
), "cross_entropy and fused_linear_cross_entropy cannot both be True."
from transformers.models.mllama import modeling_mllama
from transformers.models.mllama.modeling_mllama import (
MllamaForCausalLM,
MllamaForConditionalGeneration,
MllamaTextModel,
)
from liger_kernel.transformers.model.mllama import lce_forward as mllama_lce_forward
if rope:
modeling_mllama.apply_rotary_pos_emb = liger_rotary_pos_emb
if rms_norm:
modeling_mllama.MllamaTextRMSNorm = LigerRMSNorm
if swiglu:
modeling_mllama.MllamaTextMLP = LigerSwiGLUMLP
if cross_entropy:
modeling_mllama.CrossEntropyLoss = LigerCrossEntropyLoss
if fused_linear_cross_entropy:
# MllamaForConditionalGeneration uses MllamaForCausalLM under the hood
# for the loss calculation, so we need to patch the forward method of MllamaForCausalLM
modeling_mllama.MllamaForCausalLM.forward = mllama_lce_forward
if model is not None:
# The model instance already exists, so we need to additionally patch the
# instance variables that reference already-instantiated modules
if isinstance(model, MllamaForConditionalGeneration):
language_model: MllamaForCausalLM = model.language_model
text_model: MllamaTextModel = language_model.model
elif isinstance(model, MllamaForCausalLM):
text_model = model.model
elif isinstance(model, MllamaTextModel):
text_model = model
else:
raise ValueError(f"Unsupported Mllama model type: {type(model)}")
if text_model:
if rms_norm:
_patch_rms_norm_module(text_model.norm)
for decoder_layer in text_model.layers:
if swiglu:
_bind_method_to_module(
decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward
)
if rms_norm:
_patch_rms_norm_module(decoder_layer.input_layernorm)
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
def apply_liger_kernel_to_mistral(
rope: bool = True,
cross_entropy: bool = False,
@@ -465,7 +543,7 @@ def apply_liger_kernel_to_qwen2_vl(
) -> None:
"""
Apply Liger kernels to replace original implementation in HuggingFace Qwen2-VL models.
NOTE: Qwen2-VL is not available in transformers<=4.44.2
NOTE: Qwen2-VL is not available in transformers<4.45.0
Args:
cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
@@ -602,6 +680,8 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
"gemma": apply_liger_kernel_to_gemma,
"gemma2": apply_liger_kernel_to_gemma2,
"llama": apply_liger_kernel_to_llama,
"mllama": apply_liger_kernel_to_mllama,
"mllama_text_model": apply_liger_kernel_to_mllama,
"mistral": apply_liger_kernel_to_mistral,
"mixtral": apply_liger_kernel_to_mixtral,
"qwen2": apply_liger_kernel_to_qwen2,
test/convergence/test_mini_models_multimodal.py
+182 -12
View File
@@ -1,34 +1,63 @@
import functools
import os
from test.utils import (
FAKE_CONFIGS_PATH,
UNTOKENIZED_DATASET_PATH,
MiniModelConfig,
assert_verbose_allclose,
load_tokenizer_config,
multimodal_collate_fn,
revert_liger_kernel_to_mllama,
revert_liger_kernel_to_qwen2_vl,
set_seed,
supports_bfloat16,
train_bpe_tokenizer,
)
import pytest
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers.models.auto.processing_auto import AutoProcessor
from transformers import PreTrainedTokenizerFast
from liger_kernel.transformers import apply_liger_kernel_to_qwen2_vl
from liger_kernel.transformers import (
apply_liger_kernel_to_mllama,
apply_liger_kernel_to_qwen2_vl,
)
try:
# Qwen2-VL is only available in transformers>4.44.2
# Qwen2-VL is only available in transformers>=4.45.0
from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
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,
)
from transformers.models.qwen2_vl.processing_qwen2_vl import Qwen2VLProcessor
QWEN2_VL_AVAILABLE = True
except ImportError:
QWEN2_VL_AVAILABLE = False
try:
# Mllama is only available in transformers>=4.45.0
from transformers.models.mllama.configuration_mllama import (
MllamaConfig,
MllamaTextConfig,
MllamaVisionConfig,
)
from transformers.models.mllama.image_processing_mllama import MllamaImageProcessor
from transformers.models.mllama.modeling_mllama import (
MllamaForConditionalGeneration,
)
from transformers.models.mllama.processing_mllama import MllamaProcessor
MLLAMA_AVAILABLE = True
except ImportError:
MLLAMA_AVAILABLE = False
torch.use_deterministic_algorithms(True)
# Only setting torch.use_deterministic_algorithms(True) throws the following error:
@@ -43,6 +72,64 @@ TEST_IMAGE_DIM = 64
MINI_MODEL_SETUPS = {}
if MLLAMA_AVAILABLE:
MINI_MODEL_SETUPS["mini_mllama"] = MiniModelConfig(
liger_kernel_patch_func=functools.partial(
apply_liger_kernel_to_mllama, fused_linear_cross_entropy=False
),
liger_kernel_patch_revert_func=revert_liger_kernel_to_mllama,
model_class=MllamaForConditionalGeneration,
mini_model_config=MllamaConfig(
vision_config=MllamaVisionConfig(
hidden_act="gelu",
hidden_size=512, # 1280
image_size=560, # 560
initializer_range=0.02,
intermediate_layers_indices=[2], # [3, 7, 15, etc...]
intermediate_size=2048, # 5120
max_num_tiles=1, # 4
norm_eps=1e-5,
num_attention_heads=4, # 16
num_channels=3,
num_global_layers=2, # 8
num_hidden_layers=8, # 32
patch_size=140, # 14
supported_aspect_ratios=[[1, 1]], # [[1, 1], [1, 2], etc... ]
vision_output_dim=1024, # 7680
),
text_config=MllamaTextConfig(
bos_token_id=0,
eos_token_id=0,
pad_token_id=0,
cross_attention_layers=[2], # [3, 8, 13, 18, etc...]
dropout=0,
hidden_act="silu",
hidden_size=1024, # 4096
initializer_range=0.02,
intermediate_size=2048, # 14336
max_position_embeddings=131_072,
num_attention_heads=8, # 32
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
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,
),
image_token_index=1, # NOTE: outside the vocab size
attn_implementation="sdpa",
),
)
if QWEN2_VL_AVAILABLE:
MINI_MODEL_SETUPS["mini_qwen2_vl"] = MiniModelConfig(
liger_kernel_patch_func=functools.partial(
@@ -54,12 +141,12 @@ if QWEN2_VL_AVAILABLE:
attention_dropout=0.0,
# Token Ids and vocab size must match those in the tokenizer/processor
# https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct/blob/main/config.json
bos_token_id=151643,
eos_token_id=151645,
vision_start_token_id=151652,
vision_end_token_id=151653,
vision_token_id=151654,
image_token_id=151655,
bos_token_id=0,
eos_token_id=0,
vision_start_token_id=1,
vision_end_token_id=2,
vision_token_id=3,
image_token_id=4,
hidden_act="silu",
hidden_size=1024, # 8192
initializer_range=0.02,
@@ -78,7 +165,7 @@ if QWEN2_VL_AVAILABLE:
sliding_window=4096,
tie_word_embeddings=True,
use_cache=False, # True
vocab_size=152064,
vocab_size=32000, # 152064,
use_sliding_window=False,
vision_config={
"depth": 4, # 32
@@ -95,7 +182,51 @@ if QWEN2_VL_AVAILABLE:
def create_processor(model_name):
if model_name == "mini_qwen2_vl":
return AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
tokenizer_config = load_tokenizer_config(
os.path.join(
FAKE_CONFIGS_PATH, "Qwen/Qwen2-VL-7B-Instruct/tokenizer_config.json"
)
)
tokenizer_base = train_bpe_tokenizer(
[
token.content
for key, token in sorted(
tokenizer_config["added_tokens_decoder"].items(),
key=lambda x: int(x[0]),
)
]
)
qwen_tokenizer = Qwen2TokenizerFast(
tokenizer_object=tokenizer_base, **tokenizer_config
)
image_processor = Qwen2VLImageProcessor()
return Qwen2VLProcessor(
image_processor=image_processor, tokenizer=qwen_tokenizer
)
elif model_name == "mini_mllama":
tokenizer_config = load_tokenizer_config(
os.path.join(
FAKE_CONFIGS_PATH,
"meta-llama/Llama-3.2-11B-Vision-Instruct/tokenizer_config.json",
)
)
tokenizer_base = train_bpe_tokenizer(
[
token.content
for key, token in sorted(
tokenizer_config["added_tokens_decoder"].items(),
key=lambda x: int(x[0]),
)
]
)
fast_tokenizer = PreTrainedTokenizerFast(
tokenizer_object=tokenizer_base, **tokenizer_config
)
image_processor = MllamaImageProcessor(size={"height": 560, "width": 560})
return MllamaProcessor(
image_processor=image_processor, tokenizer=fast_tokenizer
)
else:
raise ValueError(f"Processor not available for model {model_name}")
@@ -129,7 +260,9 @@ def create_multimodal_dataset(model_name: str):
"content": [{"type": "text", "text": example["text"]}],
},
]
example["text"] = processor.apply_chat_template(conversation, tokenize=False)
example["text"] = processor.tokenizer.apply_chat_template(
conversation, tokenize=False
)
return example
def preprocess_function(examples):
@@ -266,6 +399,43 @@ def run_mini_model_multimodal(
),
],
),
pytest.param(
"mini_mllama",
32,
1e-4,
torch.float32,
1e-8,
1e-5,
5e-3,
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
not MLLAMA_AVAILABLE,
reason="Mllama not available in this version of transformers",
),
),
pytest.param(
"mini_mllama",
32,
1e-4,
torch.bfloat16,
1e-3,
1e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(
not supports_bfloat16(), reason="bfloat16 not supported on this GPU"
),
pytest.mark.skipif(
not MLLAMA_AVAILABLE,
reason="Mllama not available in this version of transformers",
),
],
),
],
)
def test_mini_model_multimodal(
test/convergence/test_mini_models_no_logits.py
+83
View File
@@ -7,6 +7,7 @@ from test.utils import (
revert_liger_kernel_to_llama,
revert_liger_kernel_to_mistral,
revert_liger_kernel_to_mixtral,
revert_liger_kernel_to_mllama,
revert_liger_kernel_to_phi3,
revert_liger_kernel_to_qwen2,
revert_liger_kernel_to_qwen2_vl,
@@ -33,11 +34,21 @@ from liger_kernel.transformers import (
apply_liger_kernel_to_llama,
apply_liger_kernel_to_mistral,
apply_liger_kernel_to_mixtral,
apply_liger_kernel_to_mllama,
apply_liger_kernel_to_phi3,
apply_liger_kernel_to_qwen2,
apply_liger_kernel_to_qwen2_vl,
)
try:
# Mllama is only available in transformers>=4.45.0
from transformers.models.mllama.configuration_mllama import MllamaTextConfig
from transformers.models.mllama.modeling_mllama import MllamaForCausalLM
MLLAMA_AVAILABLE = True
except ImportError:
MLLAMA_AVAILABLE = False
try:
# Qwen2-VL is only available in transformers>4.44.2
from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig
@@ -278,6 +289,41 @@ MINI_MODEL_SETUPS = {
),
}
if MLLAMA_AVAILABLE:
MINI_MODEL_SETUPS["mini_mllama"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_mllama,
liger_kernel_patch_revert_func=revert_liger_kernel_to_mllama,
model_class=MllamaForCausalLM,
mini_model_config=MllamaTextConfig(
bos_token_id=1, # 128000
eos_token_id=2, # 128001
pad_token_id=2,
cross_attention_layers=None,
dropout=0,
hidden_act="silu",
hidden_size=1024, # 4096
initializer_range=0.02,
intermediate_size=2048, # 14336
max_position_embeddings=131_072,
num_attention_heads=8, # 32
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
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
),
)
if QWEN2_VL_AVAILABLE:
MINI_MODEL_SETUPS["mini_qwen2_vl"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_qwen2_vl,
@@ -418,6 +464,43 @@ def run_mini_model(
not supports_bfloat16(), reason="bfloat16 not supported on this GPU"
),
),
pytest.param(
"mini_mllama",
32,
1e-4,
torch.float32,
1e-8,
1e-5,
5e-3,
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
not MLLAMA_AVAILABLE,
reason="Mllama not available in this version of transformers",
),
),
pytest.param(
"mini_mllama",
32,
1e-4,
torch.bfloat16,
1e-3,
1e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(
not supports_bfloat16(), reason="bfloat16 not supported on this GPU"
),
pytest.mark.skipif(
not MLLAMA_AVAILABLE,
reason="Mllama not available in this version of transformers",
),
],
),
("mini_qwen2", 32, 1e-4, torch.float32, 1e-8, 1e-5, 5e-3, 1e-5, 5e-3, 1e-5),
pytest.param(
"mini_qwen2",
test/resources/fake_configs/Qwen/Qwen2-VL-7B-Instruct/tokenizer_config.json
+55
View File
@@ -0,0 +1,55 @@
{
"added_tokens_decoder": {
"0": {
"content": "<|unk|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"1": {
"content": "<|vision_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"2": {
"content": "<|vision_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"3": {
"content": "<|vision_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"4": {
"content": "<|image_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"additional_special_tokens": ["<|im_start|>", "<|im_end|>", "<|object_ref_start|>","<|object_ref_end|>","<|box_start|>","<|box_end|>","<|quad_start|>","<|quad_end|>","<|vision_start|>","<|vision_end|>","<|vision_pad|>","<|image_pad|>","<|video_pad|>"],
"bos_token": null,
"chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}",
"clean_up_tokenization_spaces": false,
"eos_token": "<|im_end|>",
"padding_side": "left",
"errors": "replace",
"model_max_length": 32768,
"pad_token": "<|endoftext|>",
"split_special_tokens": false,
"unk_token": null
}
test/resources/fake_configs/meta-llama/Llama-3.2-11B-Vision-Instruct/tokenizer_config.json
+31
View File
File diff suppressed because one or more lines are too long
test/transformers/test_mm_int8int2.py
+1 -1
View File
@@ -3,8 +3,8 @@ import torch
from liger_kernel.ops.experimental.mm_int8int2 import (
matmul,
unpack_weights,
pack_weights,
unpack_weights,
)
test/transformers/test_monkey_patch.py
+129
View File
@@ -32,6 +32,7 @@ def test_import_from_root():
apply_liger_kernel_to_llama,
apply_liger_kernel_to_mistral,
apply_liger_kernel_to_mixtral,
apply_liger_kernel_to_mllama,
apply_liger_kernel_to_phi3,
apply_liger_kernel_to_qwen2,
apply_liger_kernel_to_qwen2_vl,
@@ -249,6 +250,134 @@ def test_apply_liger_kernel_to_instance_for_llama():
) == inspect.getsource(LigerRMSNorm.forward)
def test_apply_liger_kernel_to_instance_for_mllama_for_conditional_generation():
# Ensure any monkey patching is cleaned up for subsequent tests
with patch("transformers.models.mllama.modeling_mllama"):
from transformers.models.mllama.modeling_mllama import (
MllamaForConditionalGeneration,
)
# Instantiate a dummy model
config = transformers.models.mllama.configuration_mllama.MllamaConfig(
torch_dtype=torch.bfloat16,
text_config=transformers.models.mllama.configuration_mllama.MllamaTextConfig(
rms_norm_eps=1e-5,
hidden_size=32,
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",
),
),
vision_config=transformers.models.mllama.configuration_mllama.MllamaVisionConfig(
rms_norm_eps=1e-5,
hidden_size=32,
intermediate_size=64,
hidden_act="gelu",
num_hidden_layers=2,
vision_output_dim=64,
),
)
dummy_model_instance = MllamaForConditionalGeneration._from_config(config)
assert isinstance(dummy_model_instance, MllamaForConditionalGeneration)
# Check that model instance variables are not yet patched with Liger modules
assert inspect.getsource(
dummy_model_instance.language_model.model.norm.forward
) != inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.language_model.model.layers:
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(
LigerSwiGLUMLP.forward
)
assert inspect.getsource(
layer.input_layernorm.forward
) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(
layer.post_attention_layernorm.forward
) != inspect.getsource(LigerRMSNorm.forward)
# Test applying kernels to the model instance
_apply_liger_kernel_to_instance(model=dummy_model_instance)
# Check that the model's instance variables were correctly patched with Liger modules
assert inspect.getsource(
dummy_model_instance.language_model.model.norm.forward
) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.language_model.model.layers:
assert inspect.getsource(layer.mlp.forward) == inspect.getsource(
LigerSwiGLUMLP.forward
)
assert inspect.getsource(
layer.input_layernorm.forward
) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(
layer.post_attention_layernorm.forward
) == inspect.getsource(LigerRMSNorm.forward)
def test_apply_liger_kernel_to_instance_for_mllama_for_causal_lm():
# Ensure any monkey patching is cleaned up for subsequent tests
with patch("transformers.models.mllama.modeling_mllama"):
from transformers.models.mllama.modeling_mllama import MllamaForCausalLM
# Instantiate a dummy model
config = transformers.models.mllama.configuration_mllama.MllamaTextConfig(
rms_norm_eps=1e-5,
hidden_size=32,
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",
),
)
dummy_model_instance = MllamaForCausalLM._from_config(config)
assert isinstance(dummy_model_instance, MllamaForCausalLM)
# Check that model instance variables are not yet patched with Liger modules
assert not isinstance(dummy_model_instance.model.norm, LigerRMSNorm)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(
LigerSwiGLUMLP.forward
)
assert inspect.getsource(
layer.input_layernorm.forward
) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(
layer.post_attention_layernorm.forward
) != inspect.getsource(LigerRMSNorm.forward)
# Test applying kernels to the model instance
_apply_liger_kernel_to_instance(model=dummy_model_instance)
# Check that the model's instance variables were correctly patched with Liger modules
assert inspect.getsource(
dummy_model_instance.model.norm.forward
) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) == inspect.getsource(
LigerSwiGLUMLP.forward
)
assert inspect.getsource(
layer.input_layernorm.forward
) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(
layer.post_attention_layernorm.forward
) == inspect.getsource(LigerRMSNorm.forward)
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"):
test/utils.py
+55
View File
@@ -1,10 +1,15 @@
import importlib
import json
import os
import random
from dataclasses import dataclass
from typing import Any, Dict, List
import torch
from tokenizers import AddedToken, Tokenizer
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.trainers import BpeTrainer
from transformers import PretrainedConfig, PreTrainedModel
from transformers.tokenization_utils_base import BatchEncoding
@@ -110,6 +115,10 @@ UNTOKENIZED_DATASET_PATH = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "resources/tiny_shakespeare.txt"
)
FAKE_CONFIGS_PATH = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "resources/fake_configs"
)
@dataclass
class MiniModelConfig:
@@ -156,6 +165,41 @@ def multimodal_collate_fn(data: List[Dict[str, Any]]):
return BatchEncoding(batch)
def load_tokenizer_config(config_path: str) -> dict:
"""Load and process tokenizer configuration from a JSON file."""
with open(config_path) as reader:
tokenizer_config = json.load(reader)
tokenizer_config["added_tokens_decoder"] = {
k: AddedToken(**v) for k, v in tokenizer_config["added_tokens_decoder"].items()
}
return tokenizer_config
def train_bpe_tokenizer(special_tokens: List[str], unk_token: str = "<|unk|>"):
"""
Train a tokenizer using the BPE algorithm.
Parameters:
unk_token (str): The token to use for unknown tokens.
special_tokens (List[str]): A list of special tokens to use.
Returns:
Tokenizer: The trained tokenizer.
"""
# Add unk_token to special_tokens if not already present
if unk_token not in special_tokens:
special_tokens.append(unk_token)
tokenizer = Tokenizer(BPE(unk_token=unk_token))
trainer = BpeTrainer(special_tokens=special_tokens)
tokenizer.pre_tokenizer = Whitespace()
file = [UNTOKENIZED_DATASET_PATH]
tokenizer.train(file, trainer)
return tokenizer
def supports_bfloat16():
if not torch.cuda.is_available():
return False
@@ -173,6 +217,17 @@ def revert_liger_kernel_to_llama():
print("Liger kernel patches have been reverted.")
def revert_liger_kernel_to_mllama():
"""
Revert all Liger kernel patches applied to MLlama.
"""
from transformers.models.mllama import modeling_mllama
importlib.reload(modeling_mllama)
print("Liger kernel patches have been reverted.")
def revert_liger_kernel_to_mistral():
"""
Revert all Liger kernel patches applied to Mistral.