Allow option to choose between fused LCE and CE loss by connermanuel · Pull Request #704 · linkedin/Liger-Kernel

connermanuel · 2025-05-06T20:00:55Z

Summary

More context in #703
This adds a kwarg to the forward method of models that allows users to choose if they want to use fused LCE or if they want to materialize logits and use regular CE loss. Currently this decision is made implicitly using model state.

Testing Done

Ran this on a downstream transformers pipeline and measured peak memory consumption. Reduced by as much as 80% on some runs as confirmed, with no effect on quality. Still running repo tests, but shouldnt expect any surprises.

Hardware Type: H100
run make test to ensure correctness
run make checkstyle to ensure code style
run make test-convergence to ensure convergence

connermanuel · 2025-05-08T18:57:58Z

Would appreciate feedback on the interface here:

skip_logits was the natural choice for developing this because it emulates the existing logic flow. Could consider naming it something like return_logits or use_fused_kernel instead.
Added an error here when skip_logits is True but there are no labels. If you don't have labels, you're not returning loss, and it doesn't make sense to return neither loss nor logits

Not sure who to flag, maybe @shivam15s ?

Tcc0403 · 2025-05-10T20:28:27Z

Thanks for the contribution! But I'm a little bit confused the goal of this PR.

My understanding of the proposal is that it enables users to access logits while the model is in training mode. Am I understanding this correctly? How is it different from patching CE without patching FLCE?

If a user only want to retrieve logits for validation, simply setting the model to evaluation mode can materialize full logits and the model is still able to calculate CE loss if labels or shift_labels is provided.

Or is it for a flexible option to swtich between CE and FLCE in training mode? Could you elaborate on the scenarios where users can benefit from this option?

chiwanpark · 2025-05-12T07:37:00Z

@Tcc0403 I think this PR is valuable when the validation metric is just loss in SFT scenario. There is no need to materialize logits if the validation metric is just validation loss. However, Liger-Kernel currently materializes the logits even if we just compute validation loss.

#703 (as posted by the PR author) shows the efficient usage of this option.

chiwanpark · 2025-05-12T07:39:41Z

@connermanuel BTW, you need to implement this feature for qwen3_moe models that is recently merged in #706.

Tcc0403 · 2025-05-12T12:55:51Z

@chiwanpark Thanks for the explanation. I misunderstood the intent of this PR. So it's more about boardening the use of flce in eval mode, not about falling back to CE in training mode.

connermanuel · 2025-05-12T16:21:50Z

@Tcc0403 @chiwanpark Thanks both. The primary intent is to enable FLCE for validation, but I suppose a secondary benefit is enabling accessing logits during training.
I also didn't add it to llava because llava never returns logits, let me know if I should put it in anyway as something of a no-op to keep the interface the same.
Added qwen3-moe support. Is convergence testing required here?

Tcc0403 · 2025-05-14T17:14:59Z

I also didn't add it to llava because llava never returns logits, let me know if I should put it in anyway as something of a no-op to keep the interface the same.

Yes, put it in llava as well. It seems we missing validation mode support for llava, keep the interface so later PR can handle it.

Added qwen3-moe support. Is convergence testing required here?

I think we can adjust the workflow of convergence test for it. Adding an additional validation step (forward pass only) after for loop to check the following metrics:

losses in eval mode and train mode
losses when skip_logits=False in eval mode
logits when skip_logits=False in eval mode

This change should be applied to mini_models.py files only since models in mini_models_with_logits.py aren't patched with flce.

connermanuel · 2025-05-14T17:42:47Z

Yes, put it in llava as well

Sounds good

Adding an additional validation step (forward pass only) after for loop to check the following metrics:

Ok, so I just add these steps then run make test-convergence?

connermanuel · 2025-05-14T21:21:15Z

@Tcc0403 I modified the tests, can you see if this is what you meant?

Tcc0403 · 2025-05-14T23:11:17Z

        batch = next(loader_iter).to(model.device)
        optimizer.zero_grad()
+        if with_liger:
+            batch["skip_logits"] = True


We don't need to pass skip_logits in training loop because default is True for liger patched model.

Sounds good, by this reasoning I should remove skip_logits=False from validation as well

We can keep skip_logits=False in validation since we want to make sure the logits can be computed correctly.

Tcc0403 · 2025-05-14T23:22:02Z

+    return {
+        "train_loss": train_loss_list,
+        "eval_loss": eval_output.loss.item(),
+        "eval_logits": eval_output.logits,
+        "model": model,
+    }


nit: I prefer keeping existing names loss and logits to minimize the difference between files.

Sounds good, I can add eval loss to the loss list then

Tcc0403 · 2025-05-14T23:23:39Z

-    #     rtol=logits_rtol,
-    # )
+    # Compare the logits from the last step
+    if expected_output["eval_logits"] is not None and actual_output["eval_logits"] is not None:


Is this line only for llava since it's the only model that doesn't generate logits?

connermanuel · 2025-05-16T17:01:50Z

On an H100, three of the convergence tests are failing:

FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_qwen2_vl-32-0.0001-dtype7-1e-05-0.1-0.005-1e-05-0.005-1e-05]
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_qwen2_5_vl-32-0.0001-dtype8-1e-05-0.1-0.005-1e-05-0.005-1e-05]
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_granite3-32-0.0001-dtype16-1e-08-0.0001-0.005-1e-05-0.005-1e-05]

They all have similar trajectories as the reference

>           raise AssertionError("\n".join(mismatch_details))
E           AssertionError: Number of mismatched elements: 29133
E           Mismatch at index (1, 124, 6158): tensor1[(1, 124, 6158)] = -0.3988599479198456, tensor2[(1, 124, 6158)] = -0.3936949074268341
E           Mismatch at index (1, 124, 8486): tensor1[(1, 124, 8486)] = -1.195212483406067, tensor2[(1, 124, 8486)] = -1.1899968385696411
E           Mismatch at index (1, 124, 14097): tensor1[(1, 124, 14097)] = -0.02671273611485958, tensor2[(1, 124, 14097)] = -0.0215948186814785
E           Mismatch at index (1, 124, 15443): tensor1[(1, 124, 15443)] = -0.045016493648290634, tensor2[(1, 124, 15443)] = -0.050152797251939774
E           Mismatch at index (1, 124, 23038): tensor1[(1, 124, 23038)] = -1.4875845909118652, tensor2[(1, 124, 23038)] = -1.492673397064209
E           ... and 29128 more mismatched elements.

I think it's a product of introducing logits tests where there were none to begin with, rather than this PR changing logit computation behavior. No problems with loss.

Tcc0403 · 2025-05-19T05:48:11Z

@connermanuel What you showed seems like a numerical error to me. Feel free to relax logits tolerances of those failing models.

connermanuel · 2025-05-21T19:01:10Z

@Tcc0403 I have updated the logits tolerances for select models. I tried to keep them reasonable but qwen_2_5_vl in particular proved to have a good amount of divergence, though judging by some comments left on those test params that seems expected.

Tcc0403 · 2025-05-22T06:02:55Z

@Tcc0403 I have updated the logits tolerances for select models. I tried to keep them reasonable but qwen_2_5_vl in particular proved to have a good amount of divergence, though judging by some comments left on those test params that seems expected.

Mark a TODO next to qwen_25_vl logits tolerances. I'll further investigate the issue.

connermanuel · 2025-05-22T16:26:20Z

There is already a TODO beside the params for those models because the loss was already relaxed. I'll adjust that todo. I also left the original logits tolerances in a comment

connermanuel · 2025-05-23T14:16:04Z

@Tcc0403 @yundai424 ~~Can I merge this, or do you need to run the workflows first?~~
Updated one more set of tolerances which dismissed the approval, after approval am I clear to merge?

Tcc0403 · 2025-05-23T16:19:47Z

@Tcc0403 @yundai424 ~~Can I merge this, or do you need to run the workflows first?~~ Updated one more set of tolerances which dismissed the approval, after approval am I clear to merge?

Yes, merging it now. Thanks for the contribution!

## Summary `llama4.lce_forward` is the only one of the 42 forwards under `transformers/model/` that has no `skip_logits` parameter. It still decides on its own: ```python if self.training and (labels is not None or shift_labels is not None): ``` linkedin#704 added `skip_logits` precisely so a caller could override that, and swept every model that existed on 2025-05-23. llama4 arrived a month later in linkedin#740 and never picked it up; everything added since has it. Because the signature does not name it, a `skip_logits` argument lands in `**kwargs` and is silently ignored, so `skip_logits=False` during training still takes the fused path and returns `logits=None`. That is the one thing the flag exists to prevent, and the convergence harness relies on it, setting `eval_batch["skip_logits"] = False` before reading `eval_output.logits`. The change mirrors `mllama`, which has the identical `LigerForCausalLMLoss` shape. `skip_logits=None` evaluates to the previous condition, so nothing moves for callers that do not pass it. ## Testing Done New `test/transformers/test_llama4.py`. Both cases stay on the non-fused branch, so they run without a GPU: - `skip_logits=False` in training mode returns logits and a loss. On `main` this fails: the flag is ignored, the fused path runs, and on a CPU box Triton raises `RuntimeError: 0 active drivers`. On a GPU box it fails on `logits is None` instead. - `skip_logits=True` with no labels raises `ValueError`, as it does on every other model. On `main` this is `DID NOT RAISE`. - `pytest test/transformers/test_monkey_patch.py test/transformers/test_llama4.py`: 65 passed / 1 skipped (63 + the 2 new; the skip is `muse_glimmer` missing from the installed transformers). - `ruff check .` and `ruff format --check .` with the v0.14.11 pinned in `.pre-commit-config.yaml`: clean. Hardware Type: CPU. Ran on transformers 5.9.0 / torch 2.11.0. --------- Co-authored-by: Vaibhav Jindal <vaibhav.jndl@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

connermanuel changed the title ~~[WIP] Allow option to choose between fused LCE and CE loss~~ Allow option to choose between fused LCE and CE loss May 8, 2025

connermanuel marked this pull request as ready for review May 8, 2025 18:55

Tcc0403 reviewed May 14, 2025

View reviewed changes

lancerts requested a review from yundai424 May 18, 2025 14:56

connermanuel requested a review from Tcc0403 May 21, 2025 19:00

connermanuel added 3 commits May 22, 2025 12:51

add skip_logits kwarg and logic

08163a8

implement qwen3_moe and llava skip_logits

7f3bb13

modify convergence tests

476152e

Tcc0403 previously approved these changes May 23, 2025

View reviewed changes

relax logits tolerances

69a24f1

connermanuel dismissed Tcc0403’s stale review via 69a24f1 May 23, 2025 15:08

Tcc0403 approved these changes May 23, 2025

View reviewed changes

Tcc0403 merged commit 7b40ed8 into linkedin:main May 23, 2025

Tcc0403 mentioned this pull request Jun 2, 2025

[RFC] Logits numerical issues in convergence test #742

Open

8 tasks

Tcc0403 mentioned this pull request Sep 24, 2025

fix(flce): add shift_labels as eval mode loss condition #888

Merged

1 of 3 tasks

yupengtang mentioned this pull request Sep 7, 2026

fix(llama4): honour skip_logits in lce_forward #1445

Merged

Conversation

connermanuel commented May 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing Done

Uh oh!

connermanuel commented May 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Tcc0403 commented May 10, 2025

Uh oh!

chiwanpark commented May 12, 2025

Uh oh!

chiwanpark commented May 12, 2025

Uh oh!

Tcc0403 commented May 12, 2025

Uh oh!

connermanuel commented May 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Tcc0403 commented May 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

connermanuel commented May 14, 2025

Uh oh!

connermanuel commented May 14, 2025

Uh oh!

Tcc0403 May 14, 2025

Choose a reason for hiding this comment

Uh oh!

connermanuel May 14, 2025

Choose a reason for hiding this comment

Uh oh!

Tcc0403 May 14, 2025

Choose a reason for hiding this comment

Uh oh!

Tcc0403 May 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

connermanuel May 15, 2025

Choose a reason for hiding this comment

Uh oh!

Tcc0403 May 14, 2025

Choose a reason for hiding this comment

Uh oh!

connermanuel May 15, 2025

Choose a reason for hiding this comment

Uh oh!

connermanuel commented May 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Tcc0403 commented May 19, 2025

Uh oh!

connermanuel commented May 21, 2025

Uh oh!

Tcc0403 commented May 22, 2025

Uh oh!

connermanuel commented May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

connermanuel commented May 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Tcc0403 commented May 23, 2025

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

connermanuel commented May 6, 2025 •

edited

Loading

connermanuel commented May 8, 2025 •

edited

Loading

connermanuel commented May 12, 2025 •

edited

Loading

Tcc0403 commented May 14, 2025 •

edited

Loading

Tcc0403 May 14, 2025 •

edited

Loading

connermanuel commented May 16, 2025 •

edited

Loading

connermanuel commented May 22, 2025 •

edited

Loading

connermanuel commented May 23, 2025 •

edited

Loading