linkedin / linkedin/Liger-Kernel

Add Liger kernels + FLCE `token_accuracy` plumbing for Muse Glimmer (`muse_glimmer`)

Open
#1,378 2 comments 1 reaction 1 assignee View on GitHub

@buffett0323 is already working on this.

Since Aug 14, 2026.

Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

### 🚀 The feature, motivation and pitch

### Feature
Add Muse Glimmer (`model_type=muse_glimmer`, e.g. `MuseGlimmerForConditionalGeneration` in recent `transformers`) to Liger’s supported model map, with
the same pattern used for other CausalLM / VLM families:
1. **FusedLinearCrossEntropy / skip-logits loss path** for training (avoid materializing full `[B, S, V]` logits).
2. **`return_token_accuracy=True` → `outputs.token_accuracy`** on the model output (e.g. via `LigerForCausalLMLoss` + `LigerCausalLMOutputWithPast`,
as in the Qwen3 patch), so trainers that request the metric can log it without a custom fallback.
### Motivation
We train Muse Glimmer-30B with **TRL `SFTTrainer` + `use_liger_kernel=True`** (long context / FSDP2 + context parallel). TRL’s `compute_loss` does:
- set `inputs["return_token_accuracy"] = True` when `use_liger_kernel=True`
- skip the logits-based `mean_token_accuracy` path
- only read `outputs.token_accuracy`
On Muse today:
- `muse_glimmer` is **not** in `MODEL_TYPE_TO_APPLY_LIGER_FN`, so `_apply_liger_kernel_to_instance` is a no-op.
- Stock Muse forward always runs `lm_head`, returns `MuseGlimmerCausalLMOutputWithPast` **without** `token_accuracy`, and sets
`accepts_loss_kwargs=False`.
- Result: TRL warns *“liger-kernel did not return token_accuracy when requested”* and drops `mean_token_accuracy` from training logs even though
loss/backprop still work.
Liger already implements `return_token_accuracy` in fused CE ops and wires it for patched models (e.g. Qwen3). We want the same end-to-end path for
Muse so:
- long-context SFT can use true fused CE (VRAM), and
- TRL/HF trainers can log token accuracy without maintaining a Muse-specific trainer subclass.
Related context:
- TRL logging of Liger token accuracy: https://github.com/huggingface/trl/pull/4302
- TRL behavior when accuracy is missing / None: https://github.com/huggingface/trl/issues/4730
- Muse modeling lives in Hugging Face `transformers` (`MuseGlimmerForConditionalGeneration`); loss is currently `self.loss_function(logits, labels,
vocab_size, **kwargs)` after always computing logits (incl. softcap / `output_multiplier`).
### Pitch
Supporting `muse_glimmer` would make `use_liger_kernel=True` actually meaningful for a newly popular multimodal/instruct stack, and align Muse with
the contract TRL already assumes for Liger-backed SFT (fused loss + optional `token_accuracy` for metrics).

### Alternatives

1. **Trainer-side fallback (what we do now)**
Subclass TRL `SFTTrainer` and temporarily disable the Liger metrics branch so accuracy is computed from Muse’s existing `logits`. Restores logging
only; does **not** give Muse fused CE / skip-logits savings.
2. **Set `use_liger_kernel=False`**
Avoids the TRL warning and restores logits-based `mean_token_accuracy`, but forgoes any future Liger benefits and still materializes full logits
(painful at long `max_length` / large vocab).
3. **Vendor patch Muse forward locally**
Manually call `LigerForCausalLMLoss(..., return_token_accuracy=...)` and return a custom output with `token_accuracy`. Works, but duplicates what
Liger’s model-specific patches already do for other architectures and drifts from upstream `transformers`.
4. **Upgrade TRL only**
Newer TRL versions harden None-handling for missing accuracy, but they still cannot invent `token_accuracy` if the model forward never returns it.
Muse still needs Liger (or equivalent) plumbing.
Preferred upstream fix: first-class Muse support in Liger (RMSNorm / SwiGLU / RoPE as applicable + FLCE with `return_token_accuracy` on the
ConditionalGeneration forward), registered in `MODEL_TYPE_TO_APPLY_LIGER_FN`.

### Additional context

Muse Glimmer is a newly released model from Meta and we request for support

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.