Update transformers version in tests-bwd and optional-deps, increase timeout, and fix tests-bwd (#794)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Update tests-bwd and optional-deps to use transformers version 4.49.0
instead of 4.44.2. Also, increase the timeout for tests.

For this change, the following modifications are also needed:
- Mllama monkey_patch fix
- Perform llava convergence tests for `transformers>=4.52.0` as we don't
materialize logits for earlier versions
- Increase test tolerance for `granite`
- Run `qwen2_vl` and `qwen2_5_vl` tests only for `transformers>=4.52.4`

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

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

dev/modal/tests.py

@@ -14,7 +14,7 @@ app = modal.App("liger_tests", image=image)
 repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
 
 
-@app.function(gpu="A10G", image=repo, timeout=60 * 45)
+@app.function(gpu="A10G", image=repo, timeout=60 * 60)
 def liger_tests():
     import subprocess
 

dev/modal/tests_bwd.py

@@ -14,7 +14,7 @@ app = modal.App("liger_tests_bwd", image=image)
 repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
 
 
-@app.function(gpu="A10G", image=repo, timeout=60 * 30)
+@app.function(gpu="A10G", image=repo, timeout=60 * 60)
 def liger_bwd_tests():
     import subprocess
 
@@ -24,9 +24,9 @@ def liger_bwd_tests():
         shell=True,
         cwd=REMOTE_ROOT_PATH,
     )
-    # force install transformers==4.44.2
+    # force install transformers==4.49.0
     subprocess.run(
-        ["uv pip install transformers==4.44.2 --system"],
+        ["uv pip install transformers==4.49.0 --system"],
         check=True,
         shell=True,
         cwd=REMOTE_ROOT_PATH,

setup.py

@@ -31,7 +31,7 @@ def get_optional_dependencies():
     """Get optional dependency groups."""
     return {
         "dev": [
-            "transformers>=4.44.2",
+            "transformers>=4.49.0",
             "matplotlib>=3.7.2",
             "flake8>=4.0.1.1",
             "black>=24.4.2",

src/liger_kernel/transformers/monkey_patch.py

@@ -537,7 +537,10 @@ def apply_liger_kernel_to_mllama(
         if isinstance(model, MllamaForConditionalGeneration):
             language_model: MllamaForCausalLM = model.language_model
             vision_model: MllamaVisionModel = model.vision_model
-            text_model: MllamaTextModel = language_model
+            if isinstance(language_model, MllamaForCausalLM):
+                text_model: MllamaTextModel = language_model.model
+            else:
+                text_model = language_model
         elif isinstance(model, MllamaForCausalLM):
             text_model = model.model
             vision_model = None

test/convergence/bf16/test_mini_models.py

@@ -957,8 +957,8 @@ def run_mini_model(
                     reason="LLaVa not available in this version of transformers",
                 ),
                 pytest.mark.skipif(
-                    version.parse(transformers.__version__) < version.parse("4.49.0"),
-                    reason="Mistral not available in transformers<=4.49.0",
+                    version.parse(transformers.__version__) < version.parse("4.52.0"),
+                    reason="LLaVa doesn't materialize logits in transformers<=4.52.0 so we can't test it",
                 ),
             ],
         ),

test/convergence/fp32/test_mini_models.py

@@ -938,8 +938,8 @@ def run_mini_model(
                     reason="LLaVa not available in this version of transformers",
                 ),
                 pytest.mark.skipif(
-                    version.parse(transformers.__version__) < version.parse("4.49.0"),
-                    reason="Mistral not available in transformers<=4.49.0",
+                    version.parse(transformers.__version__) < version.parse("4.52.0"),
+                    reason="LLaVa doesn't materialize logits in transformers<=4.52.0 so we can't test it",
                 ),
             ],
         ),
@@ -1103,7 +1103,7 @@ def run_mini_model(
             torch.float32,
             1e-8,
             1e-4,
-            5e-3,  # 4e-3
+            4e-2,  # 4e-3
             1e-5,  # 1e-5
             5e-3,
             1e-5,

test/transformers/test_monkey_patch.py

@@ -74,24 +74,6 @@ def is_llama4_available():
         return False
 
 
-def is_qwen2_vl_available():
-    try:
-        import transformers.models.qwen2_vl  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
-def is_qwen2_5_vl_available():
-    try:
-        import transformers.models.qwen2_5_vl  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
 def is_qwen3_available():
     try:
         import transformers.models.qwen3  # noqa: F401
@@ -365,6 +347,7 @@ def test_apply_liger_kernel_to_instance_for_mllama_for_conditional_generation():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.mllama.modeling_mllama"):
         from transformers.models.mllama.modeling_mllama import MllamaForConditionalGeneration
+        from transformers.models.mllama.modeling_mllama import MllamaTextModel
 
         # Instantiate a dummy model
         config = transformers.models.mllama.configuration_mllama.MllamaConfig(
@@ -398,10 +381,14 @@ def test_apply_liger_kernel_to_instance_for_mllama_for_conditional_generation():
 
         # Check that model instance variables are not yet patched with Liger modules
         assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(mllama_lce_forward)
-        assert inspect.getsource(dummy_model_instance.language_model.norm.forward) != inspect.getsource(
-            LigerRMSNorm.forward
-        )
-        for layer in dummy_model_instance.language_model.layers:
+
+        if isinstance(dummy_model_instance.language_model, MllamaTextModel):
+            language_model = dummy_model_instance.language_model
+        else:
+            language_model = dummy_model_instance.language_model.model
+
+        assert inspect.getsource(language_model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
+        for layer in language_model.layers:
             assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward)
             assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
             assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
@@ -428,10 +415,8 @@ def test_apply_liger_kernel_to_instance_for_mllama_for_conditional_generation():
 
         # Check that the model's instance variables were correctly patched with Liger modules
         assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(mllama_lce_forward)
-        assert inspect.getsource(dummy_model_instance.language_model.norm.forward) == inspect.getsource(
-            LigerRMSNorm.forward
-        )
-        for layer in dummy_model_instance.language_model.layers:
+        assert inspect.getsource(language_model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
+        for layer in language_model.layers:
             assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward)
             assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
             assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
@@ -452,7 +437,6 @@ def test_apply_liger_kernel_to_instance_for_mllama_for_conditional_generation():
             assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(
                 LigerLayerNorm.forward
             )
-
         try:
             print(dummy_model_instance)
         except Exception as e:
@@ -1130,7 +1114,10 @@ def test_apply_liger_kernel_to_instance_for_qwen3_moe():
             pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
 
 
-@pytest.mark.skipif(not is_qwen2_vl_available(), reason="qwen2_vl module not available")
+@pytest.mark.skipif(
+    transformer_version < version.parse("4.52.4"),
+    reason="Qwen2-VL support is only compatible with transformers >= 4.52.4",
+)
 def test_apply_liger_kernel_to_instance_for_qwen2_vl_for_conditional_generation():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.qwen2_vl.modeling_qwen2_vl"):
@@ -1196,7 +1183,10 @@ def test_apply_liger_kernel_to_instance_for_qwen2_vl_for_conditional_generation(
             pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
 
 
-@pytest.mark.skipif(not is_qwen2_vl_available(), reason="qwen2_vl module not available")
+@pytest.mark.skipif(
+    transformer_version < version.parse("4.52.4"),
+    reason="Qwen2-VL support is only compatible with transformers >= 4.52.4",
+)
 def test_apply_liger_kernel_to_instance_for_qwen2_vl():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.qwen2_vl.modeling_qwen2_vl"):
@@ -1262,7 +1252,10 @@ def test_apply_liger_kernel_to_instance_for_qwen2_vl():
             pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
 
 
-@pytest.mark.skipif(not is_qwen2_vl_available(), reason="qwen2_vl module not available")
+@pytest.mark.skipif(
+    transformer_version < version.parse("4.52.4"),
+    reason="Qwen2-VL support is only compatible with transformers >= 4.52.4",
+)
 def test_apply_liger_kernel_to_instance_for_qwen2_vl_text():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.qwen2_vl.modeling_qwen2_vl"):
@@ -1310,7 +1303,10 @@ def test_apply_liger_kernel_to_instance_for_qwen2_vl_text():
             pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
 
 
-@pytest.mark.skipif(not is_qwen2_5_vl_available(), reason="qwen2_5_vl module not available")
+@pytest.mark.skipif(
+    transformer_version < version.parse("4.52.4"),
+    reason="Qwen2.5-VL support is only compatible with transformers >= 4.52.4",
+)
 def test_apply_liger_kernel_to_instance_for_qwen2_5_vl():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.qwen2_5_vl.modeling_qwen2_5_vl"):
@@ -1376,7 +1372,10 @@ def test_apply_liger_kernel_to_instance_for_qwen2_5_vl():
             pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
 
 
-@pytest.mark.skipif(not is_qwen2_5_vl_available(), reason="qwen2_5_vl module not available")
+@pytest.mark.skipif(
+    transformer_version < version.parse("4.52.4"),
+    reason="Qwen2.5-VL support is only compatible with transformers >= 4.52.4",
+)
 def test_apply_liger_kernel_to_instance_for_qwen2_5_vl_for_conditional_generation():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.qwen2_5_vl.modeling_qwen2_5_vl"):
@@ -1442,7 +1441,10 @@ def test_apply_liger_kernel_to_instance_for_qwen2_5_vl_for_conditional_generatio
             pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
 
 
-@pytest.mark.skipif(not is_qwen2_5_vl_available(), reason="qwen2_5_vl module not available")
+@pytest.mark.skipif(
+    transformer_version < version.parse("4.52.4"),
+    reason="Qwen2.5-VL support is only compatible with transformers >= 4.52.4",
+)
 def test_apply_liger_kernel_to_instance_for_qwen2_5_vl_text():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.qwen2_5_vl.modeling_qwen2_5_vl"):