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..acad48c52 --- /dev/null +++ b/test/transformers/test_qwen2_vl_forward.py @@ -0,0 +1,145 @@ +import ast +import importlib + +from pathlib import Path +from unittest.mock import patch + +import pytest + +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 _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(): + 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 + + +@pytest.mark.parametrize( + "module_name", + ( + "liger_kernel.transformers.model.qwen2_vl", + "liger_kernel.transformers.model.qwen2_5_vl", + ), +) +@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") + 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__ + + 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