Skip to content

Commit

Permalink
🦁 Update Liger-Kernels
Browse files Browse the repository at this point in the history
🦁 Update Liger-Kernels
  • Loading branch information
alekseymalakhov11 authored Nov 26, 2024
2 parents 880f77f + b7186e8 commit 809eb0e
Show file tree
Hide file tree
Showing 12 changed files with 2,173 additions and 2,445 deletions.
3,950 changes: 2,128 additions & 1,822 deletions poetry.lock

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ faiss-gpu = { version = "^1.7.2", optional = true }
deepspeed = { version = "0.12", optional = true }
accelerate = { version = "0.29.0", optional = true }
vllm = {version = "0.5.3", optional = true}
liger-kernel-nightly = {version = "^0.4.2.dev20241117192137", optional = true}


[tool.poetry.dev-dependencies]
Expand All @@ -70,7 +71,7 @@ annotated-types = "^0.6.0"

[tool.poetry.extras]
deepspeed = ["deepspeed", "accelerate"]
gpu = ["faiss-gpu", "xformers", "vllm"]
gpu = ["faiss-gpu", "xformers", "vllm", "liger-kernel-nightly"]

[build-system]
requires = ["poetry-core>=1.0.0"]
Expand Down
26 changes: 24 additions & 2 deletions turbo_alignment/common/tf/loaders/model/model.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,9 @@
import torch
from liger_kernel.transformers import (
apply_liger_kernel_to_gemma2,
apply_liger_kernel_to_llama,
apply_liger_kernel_to_qwen2,
)
from peft import PeftModel, get_peft_model, prepare_model_for_int8_training
from transformers import PreTrainedModel, PreTrainedTokenizerBase
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
Expand All @@ -7,7 +12,6 @@
PeftConfigRegistry,
TransformersAutoModelRegistry,
)
from turbo_alignment.modeling.liger_kernels import apply_liger_kernel_to_gemma2
from turbo_alignment.settings.model import (
ModelForPeftSettings,
ModelType,
Expand Down Expand Up @@ -47,10 +51,28 @@ def load_model(
tokenizer: PreTrainedTokenizerBase,
) -> PreTrainedModel:
if model_settings.liger_kernels_settings is not None:
apply_liger_kernel_to_llama(
rope=model_settings.liger_kernels_settings.use_rope,
cross_entropy=model_settings.liger_kernels_settings.use_cross_entropy,
swiglu=model_settings.liger_kernels_settings.use_mlp,
rms_norm=model_settings.liger_kernels_settings.use_rms_norm,
fused_linear_cross_entropy=model_settings.liger_kernels_settings.use_fused_linear_cross_entropy,
)

apply_liger_kernel_to_gemma2(
rope=model_settings.liger_kernels_settings.use_rope,
cross_entropy=model_settings.liger_kernels_settings.use_cross_entropy,
geglu=model_settings.liger_kernels_settings.use_geglu,
geglu=model_settings.liger_kernels_settings.use_mlp,
rms_norm=model_settings.liger_kernels_settings.use_rms_norm,
fused_linear_cross_entropy=model_settings.liger_kernels_settings.use_fused_linear_cross_entropy,
)

apply_liger_kernel_to_qwen2(
rope=model_settings.liger_kernels_settings.use_rope,
cross_entropy=model_settings.liger_kernels_settings.use_cross_entropy,
swiglu=model_settings.liger_kernels_settings.use_mlp,
rms_norm=model_settings.liger_kernels_settings.use_rms_norm,
fused_linear_cross_entropy=model_settings.liger_kernels_settings.use_fused_linear_cross_entropy,
)

model = TransformersAutoModelRegistry.by_name(model_settings.model_type).from_pretrained(
Expand Down
1 change: 1 addition & 0 deletions turbo_alignment/metrics/distinctness.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from collections import defaultdict

from transformers.tokenization_utils_base import PreTrainedTokenizerBase

from turbo_alignment.metrics.metric import Metric
Expand Down
1 change: 1 addition & 0 deletions turbo_alignment/metrics/registry.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from enum import Enum

from pydantic import field_validator

from turbo_alignment.common.registry import Registrable
Expand Down
3 changes: 0 additions & 3 deletions turbo_alignment/modeling/liger_kernels/__init__.py

This file was deleted.

199 changes: 0 additions & 199 deletions turbo_alignment/modeling/liger_kernels/cross_entropy.py

This file was deleted.

Loading

0 comments on commit 809eb0e

Please sign in to comment.