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

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR is a follow-up to #830, exposing `accum_dtype` option for monkey
patch functions.
All bf16 convergence tests related to fused linear cross entropy are
also enforced to run with `accum_dtype=torch.float32` for numerical
stability.

Related: #512, #742, #827, #850 
<!---
## 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 <[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()