Fix ci (#853)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This fixes the CI.
The mllama model does not support the accum_dtype param for transfomers
version 4.49.0.
The qwen model required out_hidden_size param in the config.
<!---
## 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

---------

Co-authored-by: Manan Shah <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>

Author: 3 people

src/liger_kernel/transformers/model/mllama.py

@@ -190,7 +190,9 @@ def lce_forward(
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
+    # Filter out accum_dtype from kwargs for model call as MllamaTextModel doesn't accept it in transformers 4.49.0
+    # but preserve it for loss function calls
+    model_kwargs = {k: v for k, v in kwargs.items() if k != "accum_dtype"}
     # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
     outputs = self.model(
         input_ids=input_ids,
@@ -206,7 +208,7 @@ def lce_forward(
         output_hidden_states=output_hidden_states,
         return_dict=return_dict,
         cache_position=cache_position,
-        **kwargs,
+        **model_kwargs,
     )
 
     hidden_states = outputs[0]

test/convergence/bf16/test_mini_models_multimodal.py

@@ -557,6 +557,7 @@
                 "hidden_size": 128,  # 1280
                 "num_heads": 16,
                 "in_chans": 3,
+                "out_hidden_size": 1024,
             },
             attn_implementation="sdpa",
         ),
@@ -860,7 +861,17 @@ def run_mini_model_multimodal(
     for i in range(num_steps):
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
-        output = model(**batch, accum_dtype=torch.float32)
+        supports_accum = getattr(model, "_supports_accum_dtype", None)
+        if supports_accum is None:
+            import inspect
+
+            params = inspect.signature(model.forward).parameters
+            supports_accum = ("accum_dtype" in params) or any(
+                p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()
+            )
+            setattr(model, "_supports_accum_dtype", supports_accum)
+
+        output = model(**batch, accum_dtype=torch.float32) if supports_accum else model(**batch)
         output.loss.backward()
         optimizer.step()
 

test/convergence/fp32/test_mini_models_multimodal.py

@@ -555,6 +555,7 @@
                 "hidden_size": 128,  # 1280
                 "num_heads": 16,
                 "in_chans": 3,
+                "out_hidden_size": 1024,
             },
             attn_implementation="sdpa",
         ),