From 2f78a9512975250a114f3c952d9fff46b481c2d7 Mon Sep 17 00:00:00 2001 From: SamJSui Date: Thu, 13 Aug 2026 14:55:09 -0500 Subject: [PATCH] Add DeepSeek V3 support --- README.md | 1 + src/liger_kernel/transformers/__init__.py | 3 + .../transformers/model/deepseek_v3.py | 81 +++++++++ src/liger_kernel/transformers/monkey_patch.py | 85 +++++++++ test/convergence/bf16/test_mini_models.py | 66 ++++++- .../bf16/test_mini_models_with_logits.py | 59 ++++++ test/convergence/fp32/test_mini_models.py | 63 ++++++- .../fp32/test_mini_models_with_logits.py | 56 ++++++ test/transformers/test_monkey_patch.py | 168 ++++++++++++++++++ test/utils.py | 12 ++ 10 files changed, 592 insertions(+), 2 deletions(-) create mode 100644 src/liger_kernel/transformers/model/deepseek_v3.py diff --git a/README.md b/README.md index 5e3fce335..6c033ffa4 100644 --- a/README.md +++ b/README.md @@ -387,6 +387,7 @@ loss.backward() | OLMo2 | `liger_kernel.transformers.apply_liger_kernel_to_olmo2` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Olmo3 | `liger_kernel.transformers.apply_liger_kernel_to_olmo3` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | GLM-4 | `liger_kernel.transformers.apply_liger_kernel_to_glm4` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | +| DeepSeek-V3 | `liger_kernel.transformers.apply_liger_kernel_to_deepseek_v3` | RMSNorm, SwiGLU (dense/shared MLPs), CrossEntropyLoss, FusedLinearCrossEntropy | | DeepSeek-V4 | `liger_kernel.transformers.apply_liger_kernel_to_deepseek_v4` | RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy | | GPT-OSS | `liger_kernel.transformers.apply_liger_kernel_to_gpt_oss` | RoPE, RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy | | InternVL3 | `liger_kernel.transformers.apply_liger_kernel_to_internvl` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index 26bdef91b..213f26afd 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -39,6 +39,7 @@ from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM # noqa: F401 from liger_kernel.transformers.monkey_patch import _apply_liger_kernel # noqa: F401 from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance # noqa: F401 + from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_deepseek_v3 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_deepseek_v4 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_exaone4 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_falcon_h1 # noqa: F401 @@ -158,6 +159,7 @@ def __getattr__(name: str): "apply_liger_kernel_to_smolvlm", "apply_liger_kernel_to_hunyuan_v1_dense", "apply_liger_kernel_to_hunyuan_v1_moe", + "apply_liger_kernel_to_deepseek_v3", "apply_liger_kernel_to_deepseek_v4", "apply_liger_kernel_to_exaone4", } @@ -249,6 +251,7 @@ def __getattr__(name: str): "apply_liger_kernel_to_smolvlm", "apply_liger_kernel_to_hunyuan_v1_dense", "apply_liger_kernel_to_hunyuan_v1_moe", + "apply_liger_kernel_to_deepseek_v3", "apply_liger_kernel_to_deepseek_v4", "apply_liger_kernel_to_exaone4", ] diff --git a/src/liger_kernel/transformers/model/deepseek_v3.py b/src/liger_kernel/transformers/model/deepseek_v3.py new file mode 100644 index 000000000..a2641ba6b --- /dev/null +++ b/src/liger_kernel/transformers/model/deepseek_v3.py @@ -0,0 +1,81 @@ +import torch + +from transformers.cache_utils import Cache +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.utils import can_return_tuple + +from liger_kernel.transformers.model.llama import lce_maybe_trainable_lm_head +from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result +from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast + + +@can_return_tuple +def lce_forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + skip_logits: bool | None = None, + **kwargs, +) -> LigerCausalLMOutputWithPast: + outputs: BaseModelOutputWithPast = 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, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + kept_hidden_states = hidden_states[:, slice_indices, :] + + shift_labels = kwargs.pop("shift_labels", None) + logits = None + loss = None + token_accuracy = None + predicted_tokens = None + + if skip_logits and labels is None and shift_labels is None: + raise ValueError("skip_logits is True, but labels and shift_labels are None") + + if skip_logits is None: + skip_logits = self.training and (labels is not None or shift_labels is not None) + + if skip_logits: + result = lce_maybe_trainable_lm_head( + self, + hidden_states=kept_hidden_states, + hidden_size=self.config.hidden_size, + labels=labels, + shift_labels=shift_labels, + **kwargs, + ) + loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result) + else: + logits = self.lm_head(kept_hidden_states) + if labels is not None or shift_labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + shift_labels=shift_labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + return LigerCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + token_accuracy=token_accuracy, + predicted_tokens=predicted_tokens, + ) diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 4d33d7e41..4ec109378 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -3377,6 +3377,90 @@ def apply_liger_kernel_to_hunyuan_v1_moe( _patch_rms_norm_module(decoder_layer.post_attention_layernorm) +def apply_liger_kernel_to_deepseek_v3( + rope: bool = False, + 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 DeepSeek-V3 models. + + NOTE: RoPE is not supported for DeepSeek-V3. Its attention uses interleaved partial RoPE, + which is incompatible with ``liger_rotary_pos_emb``. Routed experts are intentionally left + unchanged; SwiGLU is only applied to dense and shared-expert MLPs. + + Args: + rope (bool): Whether to apply Liger's rotary position embedding. Default is False. + Currently unsupported; emits a warning and is a no-op. + 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 to dense and shared-expert MLPs. Default is True. + model (PreTrainedModel): The model instance to apply Liger kernels to, if already 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.deepseek_v3 import modeling_deepseek_v3 + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3Model + + from liger_kernel.transformers.model.deepseek_v3 import lce_forward as deepseek_v3_lce_forward + from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP + + if rope: + logger.warning_once( + "rope=True is not supported for DeepSeek-V3: interleaved partial RoPE is " + "incompatible with liger_rotary_pos_emb. Skipping rope kernel swap." + ) + + if rms_norm: + modeling_deepseek_v3.DeepseekV3RMSNorm = LigerRMSNorm + + if cross_entropy: + from transformers.loss.loss_utils import nn + + nn.functional.cross_entropy = liger_cross_entropy + + if fused_linear_cross_entropy: + if model is not None: + model.forward = MethodType(deepseek_v3_lce_forward, model) + else: + modeling_deepseek_v3.DeepseekV3ForCausalLM.forward = deepseek_v3_lce_forward + + if swiglu: + modeling_deepseek_v3.DeepseekV3MLP = LigerQwen3MoeSwiGLUMLP + + if model is not None: + base_model: DeepseekV3Model = getattr(model, model.base_model_prefix, model) + + if rms_norm: + _patch_rms_norm_module(base_model.norm) + for decoder_layer in base_model.layers: + if swiglu: + shared_experts = getattr(decoder_layer.mlp, "shared_experts", None) + if shared_experts is not None: + _patch_swiglu_module(shared_experts, LigerQwen3MoeSwiGLUMLP) + elif not hasattr(decoder_layer.mlp, "experts"): + _patch_swiglu_module(decoder_layer.mlp, LigerQwen3MoeSwiGLUMLP) + if rms_norm: + _patch_rms_norm_module(decoder_layer.input_layernorm) + _patch_rms_norm_module(decoder_layer.post_attention_layernorm) + q_a_layernorm = getattr(decoder_layer.self_attn, "q_a_layernorm", None) + if q_a_layernorm is not None: + _patch_rms_norm_module(q_a_layernorm) + kv_a_layernorm = getattr(decoder_layer.self_attn, "kv_a_layernorm", None) + if kv_a_layernorm is not None: + _patch_rms_norm_module(kv_a_layernorm) + + def apply_liger_kernel_to_deepseek_v4( rope: bool = False, cross_entropy: bool = False, @@ -3532,6 +3616,7 @@ def __init__(self, hidden_size, eps=1e-6, **kwargs): # Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py MODEL_TYPE_TO_APPLY_LIGER_FN = { + "deepseek_v3": apply_liger_kernel_to_deepseek_v3, "deepseek_v4": apply_liger_kernel_to_deepseek_v4, "gemma": apply_liger_kernel_to_gemma, "gemma2": apply_liger_kernel_to_gemma2, diff --git a/test/convergence/bf16/test_mini_models.py b/test/convergence/bf16/test_mini_models.py index c78a97477..893f0dc0c 100644 --- a/test/convergence/bf16/test_mini_models.py +++ b/test/convergence/bf16/test_mini_models.py @@ -24,6 +24,7 @@ from transformers.models.qwen2 import Qwen2Config from transformers.models.qwen2 import Qwen2ForCausalLM +from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v3 from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v4 from liger_kernel.transformers import apply_liger_kernel_to_exaone4 from liger_kernel.transformers import apply_liger_kernel_to_falcon_h1 @@ -68,6 +69,7 @@ from test.utils import get_logprobs from test.utils import get_topk from test.utils import require_deterministic +from test.utils import revert_liger_kernel_to_deepseek_v3 from test.utils import revert_liger_kernel_to_deepseek_v4 from test.utils import revert_liger_kernel_to_exaone4 from test.utils import revert_liger_kernel_to_falcon_h1 @@ -340,6 +342,14 @@ except ImportError: HUNYUAN_V1_AVAILABLE = False +try: + from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + DEEPSEEK_V3_AVAILABLE = True +except ImportError: + DEEPSEEK_V3_AVAILABLE = False + try: from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4ForCausalLM @@ -1648,6 +1658,35 @@ ), ) +if DEEPSEEK_V3_AVAILABLE: + MINI_MODEL_SETUPS["mini_deepseek_v3"] = MiniModelConfig( + liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v3, + liger_kernel_patch_revert_func=revert_liger_kernel_to_deepseek_v3, + model_class=DeepseekV3ForCausalLM, + mini_model_config=DeepseekV3Config( + vocab_size=32000, + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=16, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=8, + kv_lora_rank=8, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + n_group=2, + topk_group=1, + first_k_dense_replace=1, + max_position_embeddings=128, + attn_implementation="sdpa", + ), + ) + if DEEPSEEK_V4_AVAILABLE: MINI_MODEL_SETUPS["mini_deepseek_v4"] = MiniModelConfig( liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v4, @@ -1767,7 +1806,13 @@ def run_mini_model( "rms_norm": True, } - if "glm4" in model_name or "qwen3_next" in model_name or "qwen3_5" in model_name or "deepseek_v4" in model_name: + if ( + "glm4" in model_name + or "qwen3_next" in model_name + or "qwen3_5" in model_name + or "deepseek_v3" in model_name + or "deepseek_v4" in model_name + ): kwargs["rope"] = False model_supports_layer_norm = "qwen2_vl" in model_name @@ -2467,6 +2512,25 @@ def run_mini_model( ), ], ), + pytest.param( + "mini_deepseek_v3", + 32, + 1e-5, + torch.bfloat16, + 1e-2, + 5e-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 DEEPSEEK_V3_AVAILABLE, + reason="DeepSeek-V3 not available in this version of transformers", + ), + ], + ), pytest.param( "mini_deepseek_v4", 32, diff --git a/test/convergence/bf16/test_mini_models_with_logits.py b/test/convergence/bf16/test_mini_models_with_logits.py index fa7a53ec0..43786ff31 100644 --- a/test/convergence/bf16/test_mini_models_with_logits.py +++ b/test/convergence/bf16/test_mini_models_with_logits.py @@ -24,6 +24,7 @@ from transformers.models.qwen2 import Qwen2Config from transformers.models.qwen2 import Qwen2ForCausalLM +from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v3 from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v4 from liger_kernel.transformers import apply_liger_kernel_to_exaone4 from liger_kernel.transformers import apply_liger_kernel_to_falcon_h1 @@ -66,6 +67,7 @@ from test.utils import get_logprobs from test.utils import get_topk from test.utils import require_deterministic +from test.utils import revert_liger_kernel_to_deepseek_v3 from test.utils import revert_liger_kernel_to_deepseek_v4 from test.utils import revert_liger_kernel_to_exaone4 from test.utils import revert_liger_kernel_to_falcon_h1 @@ -320,6 +322,14 @@ except ImportError: HUNYUAN_V1_AVAILABLE = False +try: + from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + DEEPSEEK_V3_AVAILABLE = True +except ImportError: + DEEPSEEK_V3_AVAILABLE = False + try: from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4ForCausalLM @@ -1551,6 +1561,35 @@ ), ) +if DEEPSEEK_V3_AVAILABLE: + MINI_MODEL_SETUPS["mini_deepseek_v3"] = MiniModelConfig( + liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v3, + liger_kernel_patch_revert_func=revert_liger_kernel_to_deepseek_v3, + model_class=DeepseekV3ForCausalLM, + mini_model_config=DeepseekV3Config( + vocab_size=32000, + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=16, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=8, + kv_lora_rank=8, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + n_group=2, + topk_group=1, + first_k_dense_replace=1, + max_position_embeddings=128, + attn_implementation="sdpa", + ), + ) + if DEEPSEEK_V4_AVAILABLE: MINI_MODEL_SETUPS["mini_deepseek_v4"] = MiniModelConfig( liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v4, @@ -1675,6 +1714,7 @@ def run_mini_model( or "llama4" in model_name or "qwen3_next" in model_name or "qwen3_5" in model_name + or "deepseek_v3" in model_name or "deepseek_v4" in model_name ): kwargs["rope"] = False @@ -2317,6 +2357,25 @@ def run_mini_model( ), ], ), + pytest.param( + "mini_deepseek_v3", + 32, + 1e-5, + torch.bfloat16, + 1e-2, + 5e-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 DEEPSEEK_V3_AVAILABLE, + reason="DeepSeek-V3 not available in this version of transformers", + ), + ], + ), pytest.param( "mini_deepseek_v4", 32, diff --git a/test/convergence/fp32/test_mini_models.py b/test/convergence/fp32/test_mini_models.py index eff183b36..cb38933a9 100644 --- a/test/convergence/fp32/test_mini_models.py +++ b/test/convergence/fp32/test_mini_models.py @@ -24,6 +24,7 @@ from transformers.models.qwen2 import Qwen2Config from transformers.models.qwen2 import Qwen2ForCausalLM +from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v3 from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v4 from liger_kernel.transformers import apply_liger_kernel_to_exaone4 from liger_kernel.transformers import apply_liger_kernel_to_falcon_h1 @@ -67,6 +68,7 @@ from test.utils import get_logprobs from test.utils import get_topk from test.utils import require_deterministic +from test.utils import revert_liger_kernel_to_deepseek_v3 from test.utils import revert_liger_kernel_to_deepseek_v4 from test.utils import revert_liger_kernel_to_exaone4 from test.utils import revert_liger_kernel_to_falcon_h1 @@ -336,6 +338,14 @@ except ImportError: HUNYUAN_V1_AVAILABLE = False +try: + from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + DEEPSEEK_V3_AVAILABLE = True +except ImportError: + DEEPSEEK_V3_AVAILABLE = False + try: from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4ForCausalLM @@ -1577,6 +1587,35 @@ ), ) +if DEEPSEEK_V3_AVAILABLE: + MINI_MODEL_SETUPS["mini_deepseek_v3"] = MiniModelConfig( + liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v3, + liger_kernel_patch_revert_func=revert_liger_kernel_to_deepseek_v3, + model_class=DeepseekV3ForCausalLM, + mini_model_config=DeepseekV3Config( + vocab_size=32000, + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=16, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=8, + kv_lora_rank=8, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + n_group=2, + topk_group=1, + first_k_dense_replace=1, + max_position_embeddings=128, + attn_implementation="sdpa", + ), + ) + if DEEPSEEK_V4_AVAILABLE: MINI_MODEL_SETUPS["mini_deepseek_v4"] = MiniModelConfig( liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v4, @@ -1696,7 +1735,13 @@ def run_mini_model( "rms_norm": True, } - if "glm4" in model_name or "qwen3_next" in model_name or "qwen3_5" in model_name or "deepseek_v4" in model_name: + if ( + "glm4" in model_name + or "qwen3_next" in model_name + or "qwen3_5" in model_name + or "deepseek_v3" in model_name + or "deepseek_v4" in model_name + ): kwargs["rope"] = False model_supports_layer_norm = "qwen2_vl" in model_name @@ -2235,6 +2280,22 @@ def run_mini_model( reason="Hunyuan_v1_moe not available in this version of transformers", ), ), + pytest.param( + "mini_deepseek_v3", + 32, + 1e-5, + torch.float32, + 1e-4, + 1e-4, + 1e-3, + 1e-4, + 1e-4, + 1e-4, + marks=pytest.mark.skipif( + not DEEPSEEK_V3_AVAILABLE, + reason="DeepSeek-V3 not available in this version of transformers", + ), + ), pytest.param( "mini_deepseek_v4", 32, diff --git a/test/convergence/fp32/test_mini_models_with_logits.py b/test/convergence/fp32/test_mini_models_with_logits.py index e6b9ba3e1..ef7dc07b8 100644 --- a/test/convergence/fp32/test_mini_models_with_logits.py +++ b/test/convergence/fp32/test_mini_models_with_logits.py @@ -24,6 +24,7 @@ from transformers.models.qwen2 import Qwen2Config from transformers.models.qwen2 import Qwen2ForCausalLM +from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v3 from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v4 from liger_kernel.transformers import apply_liger_kernel_to_exaone4 from liger_kernel.transformers import apply_liger_kernel_to_falcon_h1 @@ -65,6 +66,7 @@ from test.utils import get_logprobs from test.utils import get_topk from test.utils import require_deterministic +from test.utils import revert_liger_kernel_to_deepseek_v3 from test.utils import revert_liger_kernel_to_deepseek_v4 from test.utils import revert_liger_kernel_to_exaone4 from test.utils import revert_liger_kernel_to_falcon_h1 @@ -329,6 +331,14 @@ except ImportError: HUNYUAN_V1_AVAILABLE = False +try: + from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + DEEPSEEK_V3_AVAILABLE = True +except ImportError: + DEEPSEEK_V3_AVAILABLE = False + try: from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4ForCausalLM @@ -1504,6 +1514,35 @@ ), ) +if DEEPSEEK_V3_AVAILABLE: + MINI_MODEL_SETUPS["mini_deepseek_v3"] = MiniModelConfig( + liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v3, + liger_kernel_patch_revert_func=revert_liger_kernel_to_deepseek_v3, + model_class=DeepseekV3ForCausalLM, + mini_model_config=DeepseekV3Config( + vocab_size=32000, + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=16, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=8, + kv_lora_rank=8, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + n_group=2, + topk_group=1, + first_k_dense_replace=1, + max_position_embeddings=128, + attn_implementation="sdpa", + ), + ) + if DEEPSEEK_V4_AVAILABLE: MINI_MODEL_SETUPS["mini_deepseek_v4"] = MiniModelConfig( liger_kernel_patch_func=apply_liger_kernel_to_deepseek_v4, @@ -1628,6 +1667,7 @@ def run_mini_model( or "llama4" in model_name or "qwen3_next" in model_name or "qwen3_5" in model_name + or "deepseek_v3" in model_name or "deepseek_v4" in model_name ): kwargs["rope"] = False @@ -2103,6 +2143,22 @@ def run_mini_model( ), ], ), + pytest.param( + "mini_deepseek_v3", + 32, + 1e-5, + torch.float32, + 1e-4, + 1e-4, + 1e-3, + 1e-4, + 1e-4, + 1e-4, + marks=pytest.mark.skipif( + not DEEPSEEK_V3_AVAILABLE, + reason="DeepSeek-V3 not available in this version of transformers", + ), + ), pytest.param( "mini_deepseek_v4", 32, diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 25099f7be..fcf8d22ab 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -25,6 +25,7 @@ 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.deepseek_v3 import lce_forward as deepseek_v3_lce_forward 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 @@ -210,6 +211,15 @@ def is_paligemma_available(): return False +def is_deepseek_v3_available(): + try: + import transformers.models.deepseek_v3 # noqa: F401 + + return True + except ImportError: + return False + + def is_deepseek_v4_available(): try: import transformers.models.deepseek_v4 # noqa: F401 @@ -285,6 +295,7 @@ def is_nemotron_available(): def test_import_from_root(): try: from liger_kernel.transformers import AutoLigerKernelForCausalLM # noqa: F401 + from liger_kernel.transformers import apply_liger_kernel_to_deepseek_v3 # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_gemma # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_gemma2 # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_gemma3 # noqa: F401 @@ -1610,6 +1621,163 @@ def test_apply_liger_kernel_to_instance_for_mixtral(): pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") +@pytest.mark.skipif(not is_deepseek_v3_available(), reason="deepseek_v3 module not available") +def test_apply_liger_kernel_to_deepseek_v3_does_not_patch_rope(): + from transformers.models.deepseek_v3 import modeling_deepseek_v3 + + original_rope_source = inspect.getsource(modeling_deepseek_v3.apply_rotary_pos_emb) + original_interleaved_rope_source = inspect.getsource(modeling_deepseek_v3.apply_rotary_pos_emb_interleave) + + monkey_patch.apply_liger_kernel_to_deepseek_v3( + rope=True, + cross_entropy=False, + fused_linear_cross_entropy=False, + rms_norm=False, + swiglu=False, + ) + + assert inspect.getsource(modeling_deepseek_v3.apply_rotary_pos_emb) == original_rope_source + assert inspect.getsource(modeling_deepseek_v3.apply_rotary_pos_emb_interleave) == original_interleaved_rope_source + + +@pytest.mark.skipif(not is_deepseek_v3_available(), reason="deepseek_v3 module not available") +@pytest.mark.parametrize("q_lora_rank", [8, None]) +def test_apply_liger_kernel_to_instance_for_deepseek_v3(q_lora_rank): + with patch("transformers.models.deepseek_v3.modeling_deepseek_v3"): + config = transformers.models.deepseek_v3.configuration_deepseek_v3.DeepseekV3Config( + vocab_size=1024, + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=16, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=q_lora_rank, + kv_lora_rank=8, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + n_group=2, + topk_group=1, + first_k_dense_replace=1, + max_position_embeddings=128, + ) + dummy_model_instance = AutoModelForCausalLM.from_config(config) + + routed_expert_sources = {} + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(deepseek_v3_lce_forward) + assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + for layer_index, layer in enumerate(dummy_model_instance.model.layers): + if hasattr(layer.mlp, "shared_experts"): + assert inspect.getsource(layer.mlp.shared_experts.forward) != inspect.getsource( + LigerQwen3MoeSwiGLUMLP.forward + ) + routed_expert_sources[layer_index] = inspect.getsource(layer.mlp.experts.forward) + else: + assert inspect.getsource(layer.mlp.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) + q_a_layernorm = getattr(layer.self_attn, "q_a_layernorm", None) + if q_lora_rank is not None: + assert inspect.getsource(q_a_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) + else: + assert q_a_layernorm is None + assert inspect.getsource(layer.self_attn.kv_a_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) + + _apply_liger_kernel_to_instance(model=dummy_model_instance) + + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(deepseek_v3_lce_forward) + assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer_index, layer in enumerate(dummy_model_instance.model.layers): + if hasattr(layer.mlp, "shared_experts"): + assert inspect.getsource(layer.mlp.shared_experts.forward) == inspect.getsource( + LigerQwen3MoeSwiGLUMLP.forward + ) + assert inspect.getsource(layer.mlp.experts.forward) == routed_expert_sources[layer_index] + else: + assert inspect.getsource(layer.mlp.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) + q_a_layernorm = getattr(layer.self_attn, "q_a_layernorm", None) + if q_lora_rank is not None: + assert inspect.getsource(q_a_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) + else: + assert q_a_layernorm is None + assert inspect.getsource(layer.self_attn.kv_a_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) + + try: + print(dummy_model_instance) + except Exception as e: + pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") + + +@pytest.mark.skipif(not is_deepseek_v3_available(), reason="deepseek_v3 module not available") +def test_apply_liger_kernel_to_deepseek_v3_preserves_forward_api(): + config = transformers.models.deepseek_v3.configuration_deepseek_v3.DeepseekV3Config( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=8, + kv_lora_rank=8, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + n_group=2, + topk_group=1, + first_k_dense_replace=1, + max_position_embeddings=16, + ) + dummy_model_instance = AutoModelForCausalLM.from_config(config) + + monkey_patch.apply_liger_kernel_to_deepseek_v3( + rope=False, + cross_entropy=False, + fused_linear_cross_entropy=True, + rms_norm=False, + swiglu=False, + model=dummy_model_instance, + ) + + with torch.no_grad(): + outputs = dummy_model_instance(input_ids=torch.tensor([[1, 2]]), return_dict=False) + + assert isinstance(outputs, tuple) + assert outputs[0].shape == (1, 2, config.vocab_size) + + with torch.no_grad(): + outputs = dummy_model_instance(input_ids=torch.tensor([[1, 2]]), logits_to_keep=1) + + assert outputs.logits.shape == (1, 1, config.vocab_size) + + shift_labels = torch.tensor([[2, 3]]) + expected_loss = torch.tensor(1.0) + with patch( + "liger_kernel.transformers.model.deepseek_v3.lce_maybe_trainable_lm_head", + return_value=expected_loss, + ) as fused_loss: + outputs = dummy_model_instance( + input_ids=torch.tensor([[1, 2]]), + shift_labels=shift_labels, + skip_logits=True, + ) + + assert outputs.loss is expected_loss + assert outputs.logits is None + fused_loss.assert_called_once() + assert fused_loss.call_args.kwargs["labels"] is None + assert fused_loss.call_args.kwargs["shift_labels"] is shift_labels + + @pytest.mark.skipif(not is_deepseek_v4_available(), reason="deepseek_v4 module not available") def test_apply_liger_kernel_to_instance_for_deepseek_v4(): with patch("transformers.models.deepseek_v4.modeling_deepseek_v4"): diff --git a/test/utils.py b/test/utils.py index 52142ac71..503469f54 100644 --- a/test/utils.py +++ b/test/utils.py @@ -872,6 +872,18 @@ def revert_liger_kernel_to_hunyuan_v1_moe(model_config: MiniModelConfig): print("Liger kernel patches have been reverted.") +def revert_liger_kernel_to_deepseek_v3(model_config: MiniModelConfig): + """ + Revert all Liger kernel patches applied to DeepSeek-V3. + """ + from transformers.models.deepseek_v3 import modeling_deepseek_v3 + + importlib.reload(modeling_deepseek_v3) + model_config.model_class = modeling_deepseek_v3.DeepseekV3ForCausalLM + + print("Liger kernel patches have been reverted.") + + def revert_liger_kernel_to_deepseek_v4(model_config: MiniModelConfig): """ Revert all Liger kernel patches applied to DeepSeek-V4.