fix qwen2vl and mllama test to pass failing tests (#571)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
1. Fix Qwen2VL mrope tests. Transformers removed BC with
Qwen2VLRotaryEmbedding which is why we need this change
2. Fix mllama tests to not use cross entropy. Using Cross Entropy causes
logits mismatch which is expected as we do inplace operation in cross
entropy kernel.
3. Runs make checkstyle --- there seems to be a bug in checkstyle
workflow where unformatted checkstyle seems to pass tests. Changed the
makefile to throw error in such cases.
<!---
## 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: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

Makefile

@@ -9,8 +9,10 @@ test:
 
 # Command to run ruff for linting and formatting code
 checkstyle:
-	ruff check . --fix; ruff_check_status=$$?; \
-	ruff format .; ruff_format_status=$$?; \
+	ruff check .; ruff_check_status=$$?; \
+	ruff format --check .; ruff_format_status=$$?; \
+	ruff check . --fix; \
+	ruff format .; \
 	if [ $$ruff_check_status -ne 0 ] || [ $$ruff_format_status -ne 0 ]; then \
 		exit 1; \
 	fi

examples/lightning/training.py

@@ -158,7 +158,7 @@ class DataModule(pl.LightningDataModule):
         for i in range(len(example["question"])):
             choices = ""
             for j in range(len(example["choices"][i])):
-                choices += f"{j+1}. {example['choices'][i][j]}; "
+                choices += f"{j + 1}. {example['choices'][i][j]}; "
             s = "Below is a question and multiple choice answers, choices separated by a semicolon. Please select the best answer for the question. "
             s += f"{QUESTION}{example['question'][i]} "
             s += f"{CHOICES}{choices} "

examples/medusa/callback.py

@@ -352,9 +352,9 @@ class EfficiencyCallback(transformers.TrainerCallback):
             else:
                 return world_size
 
-        assert (
-            world_size != 0
-        ), "WORLD_SIZE should be set to a positive integer. For single GPU training, please explicitly set WORLD_SIZE=1."
+        assert world_size != 0, (
+            "WORLD_SIZE should be set to a positive integer. For single GPU training, please explicitly set WORLD_SIZE=1."
+        )
 
         # TODO: add deepspeed support
         return world_size

src/liger_kernel/ops/cross_entropy.py

@@ -289,9 +289,9 @@ def cross_entropy_forward(
     weight_sum = 0.0
     if weight is not None:
         assert weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {weight.shape}"
-        assert torch.is_floating_point(
-            weight
-        ), f"If given, weight has to be a Tensor of floating point dtype. Got: {weight.dtype}"
+        assert torch.is_floating_point(weight), (
+            f"If given, weight has to be a Tensor of floating point dtype. Got: {weight.dtype}"
+        )
         sum_non_ignore_weight = torch.gather(weight, dim=0, index=target.masked_select(target_mask)).sum().item()
         weight_sum = weight.sum().item()
         # ensure weight is contiguous

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -58,9 +58,9 @@ def fused_linear_cross_entropy_forward(
     ce_weight_sum = 0.0
     if ce_weight is not None:
         assert ce_weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {ce_weight.shape}"
-        assert torch.is_floating_point(
-            ce_weight
-        ), f"If given, weight has to be a Tensor of floating point dtype. Got: {ce_weight.dtype}"
+        assert torch.is_floating_point(ce_weight), (
+            f"If given, weight has to be a Tensor of floating point dtype. Got: {ce_weight.dtype}"
+        )
         total_sum_non_ignore_ce_weight = (
             torch.gather(ce_weight, dim=0, index=target.masked_select(target_mask)).sum().item()
         )

src/liger_kernel/ops/fused_linear_jsd.py

@@ -195,9 +195,9 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
         """
         has_label = False
         if shift_labels is not None:
-            assert shift_labels.shape == (
-                teacher_input.shape[0],
-            ), f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            assert shift_labels.shape == (teacher_input.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
             shift_labels = shift_labels.contiguous()
             has_label = True
 

src/liger_kernel/ops/jsd.py

@@ -157,9 +157,9 @@ class LigerJSDFunction(torch.autograd.Function):
         """
         has_label = False
         if shift_labels is not None:
-            assert shift_labels.shape == (
-                _input.shape[0],
-            ), f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            assert shift_labels.shape == (_input.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
             shift_labels = shift_labels.contiguous()
             has_label = True
 

src/liger_kernel/ops/layer_norm.py

@@ -147,9 +147,9 @@ def layer_norm_forward(X, W, B, eps):
     Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
     Mean = torch.empty(n_rows, dtype=X.dtype, device=X.device)
     RSTD = torch.empty(n_rows, dtype=X.dtype, device=X.device)
-    assert (
-        X.shape[1] == W.shape[0]
-    ), f"Incompatible hidden size dimension between input tensor with shape[1] = {X.shape[1]} and weight tensor with shape[0] = {W.shape[0]}"
+    assert X.shape[1] == W.shape[0], (
+        f"Incompatible hidden size dimension between input tensor with shape[1] = {X.shape[1]} and weight tensor with shape[0] = {W.shape[0]}"
+    )
 
     _layer_norm_forward_kernel[(n_rows,)](
         Y,

src/liger_kernel/ops/utils.py

@@ -49,8 +49,7 @@ def calculate_settings(n):
     BLOCK_SIZE = triton.next_power_of_2(n)
     if BLOCK_SIZE > MAX_FUSED_SIZE:
         raise RuntimeError(
-            f"Cannot launch Triton kernel since n = {n} exceeds "
-            f"the recommended Triton blocksize = {MAX_FUSED_SIZE}."
+            f"Cannot launch Triton kernel since n = {n} exceeds the recommended Triton blocksize = {MAX_FUSED_SIZE}."
         )
 
     num_warps = 4

src/liger_kernel/transformers/cross_entropy.py

@@ -17,9 +17,9 @@ class LigerCrossEntropyLoss(torch.nn.Module):
         return_z_loss: bool = False,
     ):
         super().__init__()
-        assert (label_smoothing >= 0) and (
-            label_smoothing <= 1
-        ), f"label_smoothing must be between 0.0 and 1.0. Got: {label_smoothing}"
+        assert (label_smoothing >= 0) and (label_smoothing <= 1), (
+            f"label_smoothing must be between 0.0 and 1.0. Got: {label_smoothing}"
+        )
         assert reduction in {
             "mean",
             "sum",

src/liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -17,9 +17,9 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         return_z_loss: bool = False,
     ):
         super().__init__()
-        assert (label_smoothing >= 0) and (
-            label_smoothing <= 1
-        ), f"label_smoothing must be between 0.0 and 1.0. Got: {label_smoothing}"
+        assert (label_smoothing >= 0) and (label_smoothing <= 1), (
+            f"label_smoothing must be between 0.0 and 1.0. Got: {label_smoothing}"
+        )
         assert reduction in {
             "mean",
             "sum",

src/liger_kernel/transformers/group_norm.py

@@ -21,9 +21,9 @@ class LigerGroupNorm(nn.Module):
             "zeros",
         ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
 
-        assert (
-            num_channels % num_groups == 0
-        ), f"Number of channels {num_channels} must be divisible by num_groups {num_groups}"
+        assert num_channels % num_groups == 0, (
+            f"Number of channels {num_channels} must be divisible by num_groups {num_groups}"
+        )
         self.num_channels = num_channels
         self.num_groups = num_groups
         self.eps = eps
@@ -34,9 +34,9 @@ class LigerGroupNorm(nn.Module):
     def forward(self, hidden_states):
         # hidden_states: (batch_size, num_channels, *)
         assert hidden_states.dim() >= 3, f"Input must have atleast 3 dimensions, got {hidden_states.dim()}"
-        assert (
-            hidden_states.size(1) == self.num_channels
-        ), f"Input tensor must have {self.num_channels} channels, got {hidden_states.size(1)}"
+        assert hidden_states.size(1) == self.num_channels, (
+            f"Input tensor must have {self.num_channels} channels, got {hidden_states.size(1)}"
+        )
         return LigerGroupNormFunction.apply(
             hidden_states,
             self.weight,

src/liger_kernel/transformers/monkey_patch.py

@@ -85,9 +85,9 @@ def apply_liger_kernel_to_llama(
         loaded. Default is None.
     """
 
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.llama import modeling_llama
     from transformers.models.llama.modeling_llama import LlamaModel
@@ -159,9 +159,9 @@ def apply_liger_kernel_to_mllama(
         loaded. Default is None.
     """
 
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.mllama import modeling_mllama
     from transformers.models.mllama.modeling_mllama import MllamaForCausalLM
@@ -261,9 +261,9 @@ def apply_liger_kernel_to_mistral(
         model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
         loaded. Default is None.
     """
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.mistral import modeling_mistral
     from transformers.models.mistral.modeling_mistral import MistralModel
@@ -321,9 +321,9 @@ def apply_liger_kernel_to_mixtral(
         loaded. Default is None.
     """
 
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.mixtral import modeling_mixtral
     from transformers.models.mixtral.modeling_mixtral import MixtralModel
@@ -393,9 +393,9 @@ def apply_liger_kernel_to_gemma(
         model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
         loaded. Default is None.
     """
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.gemma import modeling_gemma
     from transformers.models.gemma.modeling_gemma import GemmaModel
@@ -467,9 +467,9 @@ def apply_liger_kernel_to_gemma2(
         model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
         loaded. Default is None.
     """
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.gemma2 import modeling_gemma2
     from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
@@ -544,9 +544,9 @@ def apply_liger_kernel_to_qwen2(
         model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
         loaded. Default is None.
     """
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.qwen2 import modeling_qwen2
     from transformers.models.qwen2.modeling_qwen2 import Qwen2Model
@@ -619,9 +619,9 @@ def apply_liger_kernel_to_qwen2_vl(
         model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
         loaded. Default is None.
     """
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.qwen2_vl import modeling_qwen2_vl
     from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLModel
@@ -689,9 +689,9 @@ def apply_liger_kernel_to_phi3(
         model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
         loaded. Default is None.
     """
-    assert not (
-        cross_entropy and fused_linear_cross_entropy
-    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
 
     from transformers.models.phi3 import modeling_phi3
     from transformers.models.phi3.modeling_phi3 import Phi3Model

test/chunked_loss/test_grpo_loss.py

@@ -96,7 +96,7 @@ class TorchLMHeadGRPO(torch.nn.Module):
         per_token_loss = -(per_token_loss - self.beta * kl_div)
 
         # Apply masking and normalize
-        loss = ((per_token_loss * attention_mask).sum() / attention_mask.sum())
+        loss = (per_token_loss * attention_mask).sum() / attention_mask.sum()
 
         # Compute metrics
         metrics = (

test/convergence/bf16/test_mini_models_multimodal.py

@@ -296,7 +296,7 @@ def run_mini_model_multimodal(
         kwargs = {
             "rope": True,
             "rms_norm": True,
-            "cross_entropy": True,
+            "cross_entropy": False,
             "layer_norm": True,
         }
 

test/convergence/bf16/test_mini_models_with_logits.py

@@ -426,7 +426,7 @@ def run_mini_model(
             kwargs["swiglu"] = True
 
         kwargs["fused_linear_cross_entropy"] = False
-        kwargs["cross_entropy"] = True
+        kwargs["cross_entropy"] = False
 
         MINI_MODEL_SETUPS[model_name].liger_kernel_patch_func(**kwargs)
     else:
@@ -634,8 +634,6 @@ def test_mini_model(
         rtol=loss_rtol,
     )
 
-    # No logits are materialized
-    # import pdb; pdb.set_trace()
     # Compare the logits from the last step
     assert_verbose_allclose(
         expected_output["logits"],

test/convergence/fp32/test_mini_models_multimodal.py

@@ -295,7 +295,7 @@ def run_mini_model_multimodal(
         kwargs = {
             "rope": True,
             "rms_norm": True,
-            "cross_entropy": True,
+            "cross_entropy": False,
             "layer_norm": True,
         }
 
@@ -327,7 +327,6 @@ def run_mini_model_multimodal(
         print(f"Step {i}, Loss: {output.loss.item()}")
         loss_list.append(output.loss.item())
 
-    MINI_MODEL_SETUPS[model_name].liger_kernel_patch_revert_func(**revert_kwargs)
     return {"loss": loss_list, "logits": output.logits, "model": model}
 
 

test/convergence/fp32/test_mini_models_with_logits.py

@@ -425,7 +425,7 @@ def run_mini_model(
             kwargs["swiglu"] = True
 
         kwargs["fused_linear_cross_entropy"] = False
-        kwargs["cross_entropy"] = True
+        kwargs["cross_entropy"] = False
 
         MINI_MODEL_SETUPS[model_name].liger_kernel_patch_func(**kwargs)
     else:

test/transformers/test_monkey_patch.py

@@ -216,9 +216,9 @@ def test_patching_apis_support_patching_model_instance():
     for func in patching_functions:
         sig = inspect.signature(func)
         # Ensure 'model' is in the parameters
-        assert (
-            "model" in sig.parameters
-        ), f"{func.__name__} does not have 'model' as an argument. All patching methods must support patching an existing model instance."
+        assert "model" in sig.parameters, (
+            f"{func.__name__} does not have 'model' as an argument. All patching methods must support patching an existing model instance."
+        )
 
 
 def test_apply_liger_kernel_to_instance_for_llama():

test/transformers/test_qwen2vl_mrope.py

@@ -4,6 +4,7 @@ import torch
 from test.utils import supports_bfloat16
 
 try:
+    from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig
     from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLRotaryEmbedding
     from transformers.models.qwen2_vl.modeling_qwen2_vl import apply_multimodal_rotary_pos_emb
 
@@ -44,7 +45,7 @@ device = infer_device()
     ],
 )
 def test_correctness(bsz, seq_len, num_q_heads, num_kv_heads, head_dim, mrope_section, dtype, atol, rtol):
-    rotary_emb = Qwen2VLRotaryEmbedding(head_dim, device=device)
+    rotary_emb = Qwen2VLRotaryEmbedding(config=Qwen2VLConfig(head_dim=head_dim), device=device)
 
     _tensor_q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
 
@@ -104,7 +105,7 @@ def test_functional_correctness(bsz, seq_len, num_q_heads, num_kv_heads, head_di
     k1 = _k.clone().requires_grad_(True)
     k2 = _k.clone().requires_grad_(True)
 
-    rotary_emb = Qwen2VLRotaryEmbedding(head_dim, device=device)
+    rotary_emb = Qwen2VLRotaryEmbedding(config=Qwen2VLConfig(head_dim=head_dim), device=device)
 
     pos_ids = torch.arange(seq_len * 3 * bsz, device=device, dtype=torch.long).view(3, bsz, seq_len)
     cos, sin = rotary_emb(k1, pos_ids)

test/triton/test_triton_monkey_patch.py

@@ -24,6 +24,6 @@ def test_import_custom_cache_manager():
     cache_manager = get_cache_manager(key=random_hex_key)
     from liger_kernel.triton.monkey_patch import LigerTritonFileCacheManager
 
-    assert isinstance(
-        cache_manager, LigerTritonFileCacheManager
-    ), "Cache manager should have been LigerTritonFileCacheManager"
+    assert isinstance(cache_manager, LigerTritonFileCacheManager), (
+        "Cache manager should have been LigerTritonFileCacheManager"
+    )