linkedin / linkedin/Liger-Kernel
Add Liger kernels + FLCE `token_accuracy` plumbing for Muse Glimmer (`muse_glimmer`)
@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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.