Add patches for Qwen3_5ForConditionalGeneration to support multimodal. (#1150)

## Summary
Add support for Qwen3.5 multimodal patches as follow-up for #1123.
Implement extra `lce_forward` for `Qwen3_5ForConditionalGeneration` and
add corresponding tests.

### Details

Modified files:
- `model/qwen3_5.py`: add `lce_forward_for_multimodal` for
`Qwen3_5ForConditionalGeneration`.
- `model/output_classes.py`: add `LigerQwen3_5CausalLMOutputWithPast` as
the return type for qwen3_5.
- `test/utils.py`: `revert_liger_kernel_to_qwen3_5` to support
`conditional_generation` for qwen3_5.
- `test_monkey_patch.py`: add
`test_apply_liger_kernel_to_instance_for_qwen3_5_for_conditional_generation`.
- `bf16/test_mini_models_multimodal.py`: add
`MINI_MODEL_SETUPS["mini_qwen3_5"]` to convergence test.
- `fp32/test_mini_models_multimodal.py`: skipped

## Testing Done

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

Author: wyt2000

src/liger_kernel/transformers/model/output_classes.py

@@ -73,6 +73,13 @@
 except Exception:
     _Qwen3VLMoeCausalLMOutputWithPast = None
 
+try:
+    from transformers.models.qwen3_5.modeling_qwen3_5 import (
+        Qwen3_5CausalLMOutputWithPast as _Qwen3_5CausalLMOutputWithPast,
+    )
+except Exception:
+    _Qwen3_5CausalLMOutputWithPast = None
+
 
 @dataclass
 class LigerCausalLMOutputWithPast(CausalLMOutputWithPast):
@@ -156,3 +163,11 @@ class LigerQwen3VLCausalLMOutputWithPast(_Qwen3VLCausalLMOutputWithPast):
     class LigerQwen3VLMoeCausalLMOutputWithPast(_Qwen3VLMoeCausalLMOutputWithPast):
         token_accuracy: Optional[torch.FloatTensor] = None
         predicted_tokens: Optional[torch.LongTensor] = None
+
+
+if _Qwen3_5CausalLMOutputWithPast is not None:
+
+    @dataclass
+    class LigerQwen3_5CausalLMOutputWithPast(_Qwen3_5CausalLMOutputWithPast):
+        token_accuracy: Optional[torch.FloatTensor] = None
+        predicted_tokens: Optional[torch.LongTensor] = None

src/liger_kernel/transformers/model/qwen3_5.py

@@ -7,6 +7,7 @@
 from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
 from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast
+from liger_kernel.transformers.model.output_classes import LigerQwen3_5CausalLMOutputWithPast
 
 
 def lce_forward(
@@ -119,3 +120,137 @@ def lce_forward(
         token_accuracy=token_accuracy,
         predicted_tokens=predicted_tokens,
     )
+
+
+def lce_forward_for_multimodal(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    pixel_values: Optional[torch.Tensor] = None,
+    pixel_values_videos: Optional[torch.FloatTensor] = None,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    mm_token_type_ids: Optional[torch.IntTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
+    **kwargs,
+) -> Union[tuple, LigerQwen3_5CausalLMOutputWithPast]:
+    r"""
+    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]`.
+    image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+        The temporal, height and width of feature shape of each image in LLM.
+    video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+        The temporal, height and width of feature shape of each video in LLM.
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
+
+    >>> model = Qwen3_5ForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-8B-Instruct")
+    >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-8B-Instruct")
+
+    >>> messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image",
+                    "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg",
+                },
+                {"type": "text", "text": "Describe the image."},
+            ],
+        }
+    ]
+
+    >>> inputs = processor.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_dict=True,
+        return_tensors="pt"
+    )
+
+    >>> # Generate
+    >>> generated_ids = model.generate(**inputs, max_new_tokens=1024)
+    >>> generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    >>> output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    >>> print(output_text)
+    ```
+    """
+    return_dict = kwargs.pop("return_dict", None)
+    if return_dict is None:
+        return_dict = self.config.use_return_dict
+
+    outputs = self.model(
+        input_ids=input_ids,
+        pixel_values=pixel_values,
+        pixel_values_videos=pixel_values_videos,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=video_grid_thw,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        mm_token_type_ids=mm_token_type_ids,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    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 is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
+        result = LigerForCausalLMLoss(
+            hidden_states=kept_hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            shift_labels=shift_labels,
+            hidden_size=self.config.text_config.hidden_size,
+            **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.text_config.vocab_size,
+                **kwargs,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        output = ((loss,) + output) if loss is not None else output
+        output = output + (token_accuracy,) if token_accuracy is not None else output
+        output = output + (predicted_tokens,) if predicted_tokens is not None else output
+        return output
+
+    return LigerQwen3_5CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        rope_deltas=outputs.rope_deltas,
+        token_accuracy=token_accuracy,
+        predicted_tokens=predicted_tokens,
+    )

src/liger_kernel/transformers/monkey_patch.py

@@ -2711,47 +2711,65 @@ def apply_liger_kernel_to_qwen3_5(
     from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM
     from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel
 
+    try:
+        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForConditionalGeneration
+    except ImportError:
+        Qwen3_5ForConditionalGeneration = None
+
     from liger_kernel.transformers.model.qwen3_5 import lce_forward as qwen3_5_lce_forward
+    from liger_kernel.transformers.model.qwen3_5 import lce_forward_for_multimodal as qwen3_5_lce_forward_for_multimodal
+    from liger_kernel.transformers.monkey_patch import _patch_rms_norm_module
+    from liger_kernel.transformers.monkey_patch import _patch_swiglu_module
     from liger_kernel.transformers.rms_norm import LigerRMSNormForQwen3Next
     from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP
 
     if rope:
         raise NotImplementedError("liger_rotary_pos_emb is not available for Qwen3_5 models.")
+
     if rms_norm:
         modeling_qwen3_5.Qwen3_5RMSNorm = LigerRMSNormForQwen3Next
+
     if cross_entropy:
         from transformers.loss.loss_utils import nn
 
+        from liger_kernel.transformers.cross_entropy import liger_cross_entropy
+
         nn.functional.cross_entropy = liger_cross_entropy
+
     if fused_linear_cross_entropy:
         if model is not None:
             if isinstance(model, Qwen3_5ForCausalLM):
                 model.forward = MethodType(qwen3_5_lce_forward, model)
+            elif isinstance(model, Qwen3_5ForConditionalGeneration):
+                model.forward = MethodType(qwen3_5_lce_forward_for_multimodal, model)
             else:
                 raise TypeError(
-                    f"fused_linear_cross_entropy is only applicable on Qwen3_5ForCausalLM. Got: {type(model)}"
+                    f"fused_linear_cross_entropy is only applicable on Qwen3_5ForCausalLM or Qwen3_5ForConditionalGeneration. Got: {type(model)}"
                 )
         else:
             modeling_qwen3_5.Qwen3_5ForCausalLM.forward = qwen3_5_lce_forward
+            if Qwen3_5ForConditionalGeneration is not None:
+                modeling_qwen3_5.Qwen3_5ForConditionalGeneration.forward = qwen3_5_lce_forward_for_multimodal
+
     if swiglu:
         modeling_qwen3_5.Qwen3_5MLP = LigerQwen3MoeSwiGLUMLP
 
     if model is not None:
         if isinstance(model, (Qwen3_5ForCausalLM, Qwen3_5TextModel)):
-            base_model: Qwen3_5TextModel = getattr(model, model.base_model_prefix, model)
+            text_model: Qwen3_5TextModel = getattr(model, model.base_model_prefix, model)
+        elif Qwen3_5ForConditionalGeneration is not None and isinstance(model, Qwen3_5ForConditionalGeneration):
+            text_model = model.model.language_model
         else:
-            raise TypeError(
-                f"Unsupported qwen3_5 model type. `model` must be `Qwen3_5ForCausalLM` or `Qwen3_5TextModel`. Got: {type(model)}"
-            )
+            raise TypeError(f"Unsupported qwen3_5 model type. Got: {type(model)}")
 
         _patch_rms_norm_module_for_qwen3_5 = partial(
             _patch_rms_norm_module, offset=1.0, casting_mode="gemma", in_place=False
         )
 
         if rms_norm:
-            _patch_rms_norm_module_for_qwen3_5(base_model.norm)
+            _patch_rms_norm_module_for_qwen3_5(text_model.norm)
 
-        for decoder_layer in base_model.layers:
+        for decoder_layer in text_model.layers:
             if rms_norm:
                 _patch_rms_norm_module_for_qwen3_5(decoder_layer.input_layernorm)
                 _patch_rms_norm_module_for_qwen3_5(decoder_layer.post_attention_layernorm)