From bb84c0d2cfe265249bc5b2269409d2060cca9cfb Mon Sep 17 00:00:00 2001 From: Shivam Sahni Date: Thu, 20 Feb 2025 11:54:47 -0800 Subject: [PATCH] fix qwen2vl and mllama test to pass failing tests (#571) ## Summary 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. ## Testing Done - Hardware Type: - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --- Makefile | 6 ++- examples/lightning/training.py | 2 +- examples/medusa/callback.py | 6 +-- src/liger_kernel/ops/cross_entropy.py | 6 +-- .../ops/fused_linear_cross_entropy.py | 6 +-- src/liger_kernel/ops/fused_linear_jsd.py | 6 +-- src/liger_kernel/ops/jsd.py | 6 +-- src/liger_kernel/ops/layer_norm.py | 6 +-- src/liger_kernel/ops/utils.py | 3 +- .../transformers/cross_entropy.py | 6 +-- .../fused_linear_cross_entropy.py | 6 +-- src/liger_kernel/transformers/group_norm.py | 12 ++--- src/liger_kernel/transformers/monkey_patch.py | 54 +++++++++---------- test/chunked_loss/test_grpo_loss.py | 2 +- .../bf16/test_mini_models_multimodal.py | 2 +- .../bf16/test_mini_models_with_logits.py | 4 +- .../fp32/test_mini_models_multimodal.py | 3 +- .../fp32/test_mini_models_with_logits.py | 2 +- test/transformers/test_monkey_patch.py | 6 +-- test/transformers/test_qwen2vl_mrope.py | 5 +- test/triton/test_triton_monkey_patch.py | 6 +-- 21 files changed, 77 insertions(+), 78 deletions(-) diff --git a/Makefile b/Makefile index 38fe22ae..bef887f8 100644 --- a/Makefile +++ b/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 diff --git a/examples/lightning/training.py b/examples/lightning/training.py index 8e58d8b1..ab116487 100644 --- a/examples/lightning/training.py +++ b/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} " diff --git a/examples/medusa/callback.py b/examples/medusa/callback.py index 33a9d194..673243b7 100644 --- a/examples/medusa/callback.py +++ b/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 diff --git a/src/liger_kernel/ops/cross_entropy.py b/src/liger_kernel/ops/cross_entropy.py index 9e4ab69e..62ee9843 100644 --- a/src/liger_kernel/ops/cross_entropy.py +++ b/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 diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 4e54db47..26776356 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/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() ) diff --git a/src/liger_kernel/ops/fused_linear_jsd.py b/src/liger_kernel/ops/fused_linear_jsd.py index f0c4d7be..55c992f0 100644 --- a/src/liger_kernel/ops/fused_linear_jsd.py +++ b/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 diff --git a/src/liger_kernel/ops/jsd.py b/src/liger_kernel/ops/jsd.py index 5b6fc521..882b4099 100644 --- a/src/liger_kernel/ops/jsd.py +++ b/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 diff --git a/src/liger_kernel/ops/layer_norm.py b/src/liger_kernel/ops/layer_norm.py index 6d527c7e..4dc97ea2 100644 --- a/src/liger_kernel/ops/layer_norm.py +++ b/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, diff --git a/src/liger_kernel/ops/utils.py b/src/liger_kernel/ops/utils.py index 8a15bf8d..736990b8 100644 --- a/src/liger_kernel/ops/utils.py +++ b/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 diff --git a/src/liger_kernel/transformers/cross_entropy.py b/src/liger_kernel/transformers/cross_entropy.py index ac522c2e..f01b1f57 100644 --- a/src/liger_kernel/transformers/cross_entropy.py +++ b/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", diff --git a/src/liger_kernel/transformers/fused_linear_cross_entropy.py b/src/liger_kernel/transformers/fused_linear_cross_entropy.py index d959a0f0..6a9f19d7 100644 --- a/src/liger_kernel/transformers/fused_linear_cross_entropy.py +++ b/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", diff --git a/src/liger_kernel/transformers/group_norm.py b/src/liger_kernel/transformers/group_norm.py index ca3d314e..22a48a84 100644 --- a/src/liger_kernel/transformers/group_norm.py +++ b/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, diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index eafce145..c23c9757 100644 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/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 diff --git a/test/chunked_loss/test_grpo_loss.py b/test/chunked_loss/test_grpo_loss.py index 712e2709..7dbf29bb 100644 --- a/test/chunked_loss/test_grpo_loss.py +++ b/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 = ( diff --git a/test/convergence/bf16/test_mini_models_multimodal.py b/test/convergence/bf16/test_mini_models_multimodal.py index 2645e26f..7661c12f 100644 --- a/test/convergence/bf16/test_mini_models_multimodal.py +++ b/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, } diff --git a/test/convergence/bf16/test_mini_models_with_logits.py b/test/convergence/bf16/test_mini_models_with_logits.py index 8605893b..19bbe7a2 100644 --- a/test/convergence/bf16/test_mini_models_with_logits.py +++ b/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"], diff --git a/test/convergence/fp32/test_mini_models_multimodal.py b/test/convergence/fp32/test_mini_models_multimodal.py index ff251b9c..fd4e30f3 100644 --- a/test/convergence/fp32/test_mini_models_multimodal.py +++ b/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} diff --git a/test/convergence/fp32/test_mini_models_with_logits.py b/test/convergence/fp32/test_mini_models_with_logits.py index 75b38874..ebeee797 100644 --- a/test/convergence/fp32/test_mini_models_with_logits.py +++ b/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: diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 811cd74c..76e937de 100644 --- a/test/transformers/test_monkey_patch.py +++ b/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(): diff --git a/test/transformers/test_qwen2vl_mrope.py b/test/transformers/test_qwen2vl_mrope.py index 1ca73b82..15ade170 100644 --- a/test/transformers/test_qwen2vl_mrope.py +++ b/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) diff --git a/test/triton/test_triton_monkey_patch.py b/test/triton/test_triton_monkey_patch.py index f031c2af..eeb8c173 100644 --- a/test/triton/test_triton_monkey_patch.py +++ b/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" + )