ci: relax DyT tol, skip MTA test for xpu, unify logit logprobs tol (#778)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Fix #761, #742.

This PR found and addressed couple of issues
- The tolerance for DyT's alpha.grad was too strict. 
	- Setting its tolerance alone and lowering the values.
- MultitokenAttention couldn't pass on xpu platform. 
	- Disabling it for now to unblock the ci.
- Llama4 in convergence test shouldn't be imported in bwd-test ci
(transformers==4.44.2).
	- Adding a import guard for Llama4.
- Fix casting issue for geglu/swiglu mlp layer.  
	- #783
- Some decprecated lce_forward functions didn't accept `skip_logits`,
causing failures for bwd-test ci.
	- Adding `skip_logits` to `lce_forward_deprecated()`s.
- Some models are no longer supported in transformers<4.49.0
- Skipping those test cases if transformers version don't meet the
requirement.
- Llama4_multimodal implicitly casts image input to `bfloat16`, couldn't
work with `float32` model.
	- Commenting llama4 out from fp32 multimodal test.

Minor change:
- Enhancing testing utility function `assert_verbose_allclose()` to show
addional information at the start of the error message.
<!---
## 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
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Signed-off-by: Tcc0403 <76503978+Tcc0403@users.noreply.github.com>

Author: Tcc0403

src/liger_kernel/ops/geglu.py

@@ -40,7 +40,7 @@ def _geglu_tanh_forward_kernel(a, b, c, stride, n_cols: tl.constexpr, BLOCK_SIZE
     tanh_arg = sqrt_2_over_pi * (a_row + 0.044715 * a_cubed)
     tanh_result = tanh(tanh_arg)
     geglu_a = 0.5 * a_row * (1 + tanh_result)
-    c_row = geglu_a * b_row
+    c_row = geglu_a.cast(b_row.dtype) * b_row
     tl.store(c + col_offsets, c_row, mask=mask)
 
 

src/liger_kernel/ops/swiglu.py

@@ -26,7 +26,7 @@ def _swiglu_forward_kernel(a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BL
     # sigmoid requires type float32
     a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
     b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
-    c_row = silu(a_row) * b_row
+    c_row = silu(a_row).cast(b_row.dtype) * b_row
     tl.store(c_ptr + col_offsets, c_row, mask=mask)
 
 

src/liger_kernel/transformers/model/gemma.py

@@ -27,6 +27,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
 
@@ -81,7 +82,14 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

src/liger_kernel/transformers/model/gemma2.py

@@ -30,6 +30,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
     **kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -85,7 +86,14 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

src/liger_kernel/transformers/model/llama.py

@@ -37,6 +37,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste llama forward but replace torch cross entropy with liger fused linear cross entropy
@@ -91,7 +92,15 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    # if in training mode, don't materialize logits
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

src/liger_kernel/transformers/model/mistral.py

@@ -133,6 +133,3 @@ def lce_forward(
         hidden_states=outputs.hidden_states,
         attentions=outputs.attentions,
     )
-
-
-# Note: Grad Acc is not fixed in mistral at transformer 4.46.1

src/liger_kernel/transformers/model/phi3.py

@@ -26,6 +26,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste phi3 forward from transfomers v4.44.2 but replace torch cross entropy with liger fused linear cross entropy
@@ -80,7 +81,14 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and labels is not None:
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

src/liger_kernel/transformers/model/qwen2.py

@@ -26,6 +26,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste Qwen2's forward but replace torch cross entropy with liger fused linear cross entropy
@@ -80,6 +81,13 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
     if self.training and (labels is not None):
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()

src/liger_kernel/transformers/monkey_patch.py

@@ -611,10 +611,17 @@ def apply_liger_kernel_to_mistral(
     if cross_entropy:
         modeling_mistral.CrossEntropyLoss = LigerCrossEntropyLoss
     if fused_linear_cross_entropy:
-        if model is not None:
-            model.forward = MethodType(mistral_lce_forward, model)
+        if transformer_version >= version.parse("4.49.0"):
+            if model is not None:
+                model.forward = MethodType(mistral_lce_forward, model)
+            else:
+                modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
         else:
-            modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
+            logger.warning(
+                "The latest version of Liger does not support transformers < 4.49.0 for llava. Please downgrade your liger version or upgrade your transformer version."
+            )
+            logger.warning("LigerFusedLinearCrossEntropy patch is not applied.")
+
     if swiglu:
         modeling_mistral.MistralMLP = LigerSwiGLUMLP