feat(FLCE): expose accum_dtype for hf model monkey patch

Signed-off-by: Tcc0403 <[email protected]>

Author: Tcc0403

src/liger_kernel/transformers/model/loss_utils.py

@@ -13,6 +13,7 @@ def fixed_fused_linear_cross_entropy(
     num_items_in_batch: Optional[int] = None,
     ignore_index: int = -100,
     final_logit_softcapping: Optional[float] = None,
+    accum_dtype: Optional[torch.dtype] = None,
     **kwargs,
 ):
     reduction = "sum" if num_items_in_batch is not None else "mean"
@@ -23,6 +24,7 @@ def fixed_fused_linear_cross_entropy(
         reduction=reduction,
         ignore_index=ignore_index,
         softcap=final_logit_softcapping,
+        accum_dtype=accum_dtype,
     )
     if reduction == "sum":
         loss = loss / num_items_in_batch

test/convergence/bf16/test_mini_models.py

@@ -926,7 +926,7 @@ def run_mini_model(
     for i in range(num_steps):
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
-        output = model(**batch)
+        output = model(**batch, accum_dtype=torch.float32)
         output.loss.backward()
         optimizer.step()
         print(f"Step {i}, Loss: {output.loss.item()}")

test/convergence/bf16/test_mini_models_multimodal.py

@@ -860,7 +860,7 @@ def run_mini_model_multimodal(
     for i in range(num_steps):
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
-        output = model(**batch)
+        output = model(**batch, accum_dtype=torch.float32)
         output.loss.backward()
         optimizer.step()