From a305a53c9d0123f8f9aecbbbb5f1504328e7e75f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=91=B8=E9=B1=BC=E5=B0=8F=E9=98=9F?= Date: Fri, 24 Jul 2026 12:19:13 +0800 Subject: [PATCH 1/3] fix: support logits_to_keep in qwen vl forwards --- .../transformers/model/qwen2_5_vl.py | 4 +- .../transformers/model/qwen2_vl.py | 4 +- test/transformers/test_qwen2_vl_forward.py | 68 +++++++++++++++++++ 3 files changed, 74 insertions(+), 2 deletions(-) create mode 100644 test/transformers/test_qwen2_vl_forward.py diff --git a/src/liger_kernel/transformers/model/qwen2_5_vl.py b/src/liger_kernel/transformers/model/qwen2_5_vl.py index ac4aae51c..6ba421515 100644 --- a/src/liger_kernel/transformers/model/qwen2_5_vl.py +++ b/src/liger_kernel/transformers/model/qwen2_5_vl.py @@ -51,6 +51,7 @@ def lce_forward( mm_token_type_ids: Optional[torch.IntTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, skip_logits: Optional[bool] = None, **kwargs, ) -> Union[Tuple, LigerQwen2_5_VLCausalLMOutputWithPast]: @@ -155,7 +156,8 @@ def lce_forward( ) loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result) else: - logits = self.lm_head(hidden_states) + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None or shift_labels is not None: diff --git a/src/liger_kernel/transformers/model/qwen2_vl.py b/src/liger_kernel/transformers/model/qwen2_vl.py index b51600a2e..c7abd18fb 100644 --- a/src/liger_kernel/transformers/model/qwen2_vl.py +++ b/src/liger_kernel/transformers/model/qwen2_vl.py @@ -50,6 +50,7 @@ def lce_forward( rope_deltas: Optional[torch.LongTensor] = None, mm_token_type_ids: Optional[torch.IntTensor] = None, cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, skip_logits: Optional[bool] = None, **kwargs, ) -> Union[Tuple, LigerQwen2VLCausalLMOutputWithPast]: @@ -151,7 +152,8 @@ def lce_forward( ) loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result) else: - logits = self.lm_head(hidden_states) + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None or shift_labels is not None: diff --git a/test/transformers/test_qwen2_vl_forward.py b/test/transformers/test_qwen2_vl_forward.py new file mode 100644 index 000000000..5db2c812a --- /dev/null +++ b/test/transformers/test_qwen2_vl_forward.py @@ -0,0 +1,68 @@ +import ast +from pathlib import Path + + +MODEL_FILES = ( + Path("src/liger_kernel/transformers/model/qwen2_vl.py"), + Path("src/liger_kernel/transformers/model/qwen2_5_vl.py"), +) + + +def _forward_function(path: Path) -> ast.FunctionDef: + tree = ast.parse(path.read_text(encoding="utf-8")) + return next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "lce_forward") + + +def _all_names(node: ast.AST, name: str) -> list[ast.AST]: + return [candidate for candidate in ast.walk(node) if isinstance(candidate, ast.Name) and candidate.id == name] + + +def test_qwen_vl_forward_declares_logits_to_keep(): + for path in MODEL_FILES: + function = _forward_function(path) + parameter_names = [argument.arg for argument in function.args.args + function.args.kwonlyargs] + assert "logits_to_keep" in parameter_names, path + + +def test_qwen_vl_forward_slices_hidden_states_before_lm_head(): + for path in MODEL_FILES: + function = _forward_function(path) + source = ast.get_source_segment(path.read_text(encoding="utf-8"), function) + assert source is not None + assert "slice(-logits_to_keep, None)" in source, path + lm_head_calls = [ + node + for node in ast.walk(function) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "lm_head" + ] + assert len(lm_head_calls) == 1, path + lm_head_argument = lm_head_calls[0].args[0] + assert isinstance(lm_head_argument, ast.Subscript), path + assert isinstance(lm_head_argument.value, ast.Name) and lm_head_argument.value.id == "hidden_states", path + + +def test_qwen_vl_forward_does_not_forward_logits_to_keep_to_base_model(): + for path in MODEL_FILES: + function = _forward_function(path) + base_model_calls = [ + node + for node in ast.walk(function) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "model" + ] + assert len(base_model_calls) == 1, path + forwarded_names = {keyword.arg for keyword in base_model_calls[0].keywords if keyword.arg is not None} + assert "logits_to_keep" not in forwarded_names, path + + +def test_qwen_vl_forward_keeps_fused_loss_on_full_hidden_states(): + for path in MODEL_FILES: + function = _forward_function(path) + source = ast.get_source_segment(path.read_text(encoding="utf-8"), function) + assert source is not None + fused_start = source.index("if skip_logits:") + logits_start = source.index("else:", fused_start) + fused_source = source[fused_start:logits_start] + assert "hidden_states=hidden_states" in fused_source, path + assert "hidden_states[:, slice_indices, :]" not in fused_source, path From 2d932a5ba6ebdcb433d4c46cc71f5c461db09285 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=91=B8=E9=B1=BC=E5=B0=8F=E9=98=9F?= Date: Fri, 24 Jul 2026 12:47:42 +0800 Subject: [PATCH 2/3] test: verify qwen vl logits selection on cpu --- test/transformers/test_qwen2_vl_forward.py | 60 ++++++++++++++++++++-- 1 file changed, 57 insertions(+), 3 deletions(-) diff --git a/test/transformers/test_qwen2_vl_forward.py b/test/transformers/test_qwen2_vl_forward.py index 5db2c812a..5df6ef886 100644 --- a/test/transformers/test_qwen2_vl_forward.py +++ b/test/transformers/test_qwen2_vl_forward.py @@ -1,6 +1,9 @@ import ast +import importlib + from pathlib import Path +import pytest MODEL_FILES = ( Path("src/liger_kernel/transformers/model/qwen2_vl.py"), @@ -33,9 +36,7 @@ def test_qwen_vl_forward_slices_hidden_states_before_lm_head(): lm_head_calls = [ node for node in ast.walk(function) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == "lm_head" + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "lm_head" ] assert len(lm_head_calls) == 1, path lm_head_argument = lm_head_calls[0].args[0] @@ -66,3 +67,56 @@ def test_qwen_vl_forward_keeps_fused_loss_on_full_hidden_states(): fused_source = source[fused_start:logits_start] assert "hidden_states=hidden_states" in fused_source, path assert "hidden_states[:, slice_indices, :]" not in fused_source, path + + +@pytest.mark.parametrize( + "module_name", + ( + "liger_kernel.transformers.model.qwen2_vl", + "liger_kernel.transformers.model.qwen2_5_vl", + ), +) +@pytest.mark.parametrize("selector_kind", ("last_two", "tensor")) +def test_qwen_vl_forward_applies_logits_to_keep_on_cpu(module_name: str, selector_kind: str): + torch = pytest.importorskip("torch") + + class DummyOutputs(tuple): + def __new__(cls, hidden_states): + output = super().__new__(cls, (hidden_states,)) + output.past_key_values = None + output.hidden_states = None + output.attentions = None + output.rope_deltas = None + return output + + class DummyBaseModel: + def __init__(self, hidden_states): + self.hidden_states = hidden_states + self.kwargs = None + + def __call__(self, **kwargs): + self.kwargs = kwargs + return DummyOutputs(self.hidden_states) + + class DummyModel: + def __init__(self, hidden_states): + self.config = type( + "Config", + (), + {"output_attentions": False, "output_hidden_states": False, "use_return_dict": False}, + )() + self.model = DummyBaseModel(hidden_states) + self.lm_head = torch.nn.Linear(hidden_states.shape[-1], 2, bias=False) + self.training = False + + hidden_states = torch.arange(12, dtype=torch.float32).reshape(1, 4, 3) + model = DummyModel(hidden_states) + selector = 2 if selector_kind == "last_two" else torch.tensor([1, 3]) + expected_indices = slice(-2, None) if selector_kind == "last_two" else selector + expected_logits = model.lm_head(hidden_states[:, expected_indices, :]) + forward = importlib.import_module(module_name).lce_forward.__wrapped__ + + outputs = forward(model, logits_to_keep=selector, return_dict=False) + + torch.testing.assert_close(outputs[0], expected_logits) + assert "logits_to_keep" not in model.model.kwargs From a2916a45785a0c7d9f10652ca432bfd87af63335 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=91=B8=E9=B1=BC=E5=B0=8F=E9=98=9F?= Date: Fri, 24 Jul 2026 13:15:41 +0800 Subject: [PATCH 3/3] test: strengthen qwen vl forward coverage --- test/transformers/test_qwen2_vl_forward.py | 123 ++++++++++++--------- 1 file changed, 73 insertions(+), 50 deletions(-) diff --git a/test/transformers/test_qwen2_vl_forward.py b/test/transformers/test_qwen2_vl_forward.py index 5df6ef886..acad48c52 100644 --- a/test/transformers/test_qwen2_vl_forward.py +++ b/test/transformers/test_qwen2_vl_forward.py @@ -2,6 +2,7 @@ import importlib from pathlib import Path +from unittest.mock import patch import pytest @@ -16,8 +17,46 @@ def _forward_function(path: Path) -> ast.FunctionDef: return next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "lce_forward") -def _all_names(node: ast.AST, name: str) -> list[ast.AST]: - return [candidate for candidate in ast.walk(node) if isinstance(candidate, ast.Name) and candidate.id == name] +def _make_dummy_model(torch): + class DummyOutputs(tuple): + def __new__(cls, hidden_states): + output = super().__new__(cls, (hidden_states,)) + output.past_key_values = None + output.hidden_states = None + output.attentions = None + output.rope_deltas = None + return output + + class DummyBaseModel: + def __init__(self, hidden_states): + self.hidden_states = hidden_states + self.kwargs = None + + def __call__(self, **kwargs): + self.kwargs = kwargs + return DummyOutputs(self.hidden_states) + + class DummyModel: + def __init__(self, hidden_states): + text_config = type("TextConfig", (), {"hidden_size": hidden_states.shape[-1], "vocab_size": 2})() + self.config = type( + "Config", + (), + { + "hidden_size": hidden_states.shape[-1], + "vocab_size": 2, + "text_config": text_config, + "output_attentions": False, + "output_hidden_states": False, + "use_return_dict": False, + }, + )() + self.model = DummyBaseModel(hidden_states) + self.lm_head = torch.nn.Linear(hidden_states.shape[-1], 2, bias=False) + self.training = False + + hidden_states = torch.arange(12, dtype=torch.float32).reshape(1, 4, 3) + return DummyModel(hidden_states), hidden_states def test_qwen_vl_forward_declares_logits_to_keep(): @@ -57,18 +96,6 @@ def test_qwen_vl_forward_does_not_forward_logits_to_keep_to_base_model(): assert "logits_to_keep" not in forwarded_names, path -def test_qwen_vl_forward_keeps_fused_loss_on_full_hidden_states(): - for path in MODEL_FILES: - function = _forward_function(path) - source = ast.get_source_segment(path.read_text(encoding="utf-8"), function) - assert source is not None - fused_start = source.index("if skip_logits:") - logits_start = source.index("else:", fused_start) - fused_source = source[fused_start:logits_start] - assert "hidden_states=hidden_states" in fused_source, path - assert "hidden_states[:, slice_indices, :]" not in fused_source, path - - @pytest.mark.parametrize( "module_name", ( @@ -76,47 +103,43 @@ def test_qwen_vl_forward_keeps_fused_loss_on_full_hidden_states(): "liger_kernel.transformers.model.qwen2_5_vl", ), ) -@pytest.mark.parametrize("selector_kind", ("last_two", "tensor")) +@pytest.mark.parametrize("selector_kind", ("all", "last_two", "tensor")) def test_qwen_vl_forward_applies_logits_to_keep_on_cpu(module_name: str, selector_kind: str): torch = pytest.importorskip("torch") - - class DummyOutputs(tuple): - def __new__(cls, hidden_states): - output = super().__new__(cls, (hidden_states,)) - output.past_key_values = None - output.hidden_states = None - output.attentions = None - output.rope_deltas = None - return output - - class DummyBaseModel: - def __init__(self, hidden_states): - self.hidden_states = hidden_states - self.kwargs = None - - def __call__(self, **kwargs): - self.kwargs = kwargs - return DummyOutputs(self.hidden_states) - - class DummyModel: - def __init__(self, hidden_states): - self.config = type( - "Config", - (), - {"output_attentions": False, "output_hidden_states": False, "use_return_dict": False}, - )() - self.model = DummyBaseModel(hidden_states) - self.lm_head = torch.nn.Linear(hidden_states.shape[-1], 2, bias=False) - self.training = False - - hidden_states = torch.arange(12, dtype=torch.float32).reshape(1, 4, 3) - model = DummyModel(hidden_states) - selector = 2 if selector_kind == "last_two" else torch.tensor([1, 3]) - expected_indices = slice(-2, None) if selector_kind == "last_two" else selector + model, hidden_states = _make_dummy_model(torch) + selector = {"all": 0, "last_two": 2, "tensor": torch.tensor([1, 3])}[selector_kind] + expected_indices = { + "all": slice(None), + "last_two": slice(-2, None), + "tensor": selector, + }[selector_kind] expected_logits = model.lm_head(hidden_states[:, expected_indices, :]) forward = importlib.import_module(module_name).lce_forward.__wrapped__ - outputs = forward(model, logits_to_keep=selector, return_dict=False) + if selector_kind == "all": + outputs = forward(model, return_dict=False) + else: + outputs = forward(model, logits_to_keep=selector, return_dict=False) torch.testing.assert_close(outputs[0], expected_logits) assert "logits_to_keep" not in model.model.kwargs + + +@pytest.mark.parametrize( + "module_name", + ( + "liger_kernel.transformers.model.qwen2_vl", + "liger_kernel.transformers.model.qwen2_5_vl", + ), +) +def test_qwen_vl_forward_keeps_fused_loss_on_full_hidden_states(module_name: str): + torch = pytest.importorskip("torch") + model, hidden_states = _make_dummy_model(torch) + labels = torch.zeros((1, hidden_states.shape[1]), dtype=torch.long) + module = importlib.import_module(module_name) + forward = module.lce_forward.__wrapped__ + + with patch.object(module, "LigerForCausalLMLoss", return_value=torch.tensor(0.0)) as fused_loss: + forward(model, labels=labels, logits_to_keep=2, skip_logits=True, return_dict=False) + + assert fused_loss.call_args.kwargs["hidden_states"] is hidden_states