Add Liger Kernel support for Muse Glimmer

[buffett0323]

Summary

Adds Liger Kernel support for Muse Glimmer (model_type=muse_glimmer), covering both capabilities requested in #1378:

1. FusedLinearCrossEntropy / skip-logits loss path, avoiding the [B, S, V] logits tensor (Muse Glimmer's vocab_size is 202,048)
2. return_token_accuracy=True → outputs.token_accuracy, so TRL's SFTTrainer can log mean_token_accuracy without a trainer-side fallback

Registered in MODEL_TYPE_TO_APPLY_LIGER_FN, so use_liger_kernel=True picks it up. Patches RoPE, RMSNorm, SwiGLU, vision LayerNorm (instance path only), CrossEntropy and FusedLinearCrossEntropy.

Details

Softcap. HF computes T * tanh(logits * m / T) where m = output_multiplier and T = final_logit_softcapping, but Liger's fused softcap only implements T * tanh(logits / T). m is folded into the hidden states instead, since (m * h) @ W.T == m * (h @ W.T) — scaling [batch, seq, hidden] rather than [vocab, hidden].

Two RMSNorm variants. MuseGlimmerRMSNorm scales by weight directly (gemma-style fp32 compute, offset 0); MuseGlimmerTextCenteredRMSNorm scales by (1 + weight) from zero-init (offset 1). The with_scale=False instances, qk_norm, embed_norm, perception_emb_norm, have no weight parameter and fall back to a torch implementation matching HF exactly; an assert guards against that fallback silently dropping a scale.

in_place=False on the TextCentered norms is required, not conservative: each decoder layer hands pre_feedforward_layernorm the same tensor it retains as the residual.

Note that token_accuracy is populated on the fused path only; with skip_logits=False (eval) it stays None, matching existing patched models such as Qwen3.

Testing Done

• test/transformers/test_muse_glimmer.py (new) — instance-path numerical parity against the unpatched model: loss plus embedding / SwiGLU / RMSNorm gradients in fp32 and bf16, the non-fused softcap branch, vision LayerNorm, and the token_accuracy / predicted_tokens contract. The convergence suite patches at module level, so this is what covers _apply_liger_kernel_to_instance — the path TRL actually takes.
• test/transformers/test_monkey_patch.py — instance patching
• Convergence in bf16 and fp32: test_mini_models, test_mini_models_multimodal, test_mini_models_with_logits

Memory

H100, bf16, forward + backward, no optimizer step. Randomly initialised (no public checkpoint), real hidden_size=6656 / vocab_size=202048, num_hidden_layers reduced from 52 so a 30B-class model fits on one GPU.

seq_len	peak mem (HF)	peak mem (Liger)	saved
1024	18.4 GB	20.0 GB	-9%
2048	19.6 GB	21.2 GB	-8%
4096	30.4 GB	23.5 GB	23%
8192	52.0 GB	28.2 GB	46%
16384	OOM	38.2 GB	—

The crossover sits near 3K: FLCE accumulates grad_weight in fp32, a fixed 5.4 GB for [202048, 6656] against bf16's 2.7 GB. Below ~3K that outweighs the logits it avoids; above it the saving compounds, and 16K only runs at all with Liger.

Throughput

FLCE trades speed for memory by design. End to end that penalty is diluted by the transformer body, and the reduced depth above overstates it:

layers	lm_head share	tok/s (HF)	tok/s (Liger)	ratio
4	41%	22,405	13,433	0.60x
8	26%	13,685	10,396	0.76x
16	15%	7,705	7,063	0.92x

lm_head share is measured from the model's parameter split, which tracks the matmul FLOP split. The shipping model's 52 layers put it at ~5%, below the last row here.

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

[buffett0323]

@vaibhavjindal FYI, thanks!

[PKUWZP]

@buffett0323 Interesting work. I am wondering if Liger kernel works out of the box for Muse Glimmer? Are you testing both Triton backend and CuteDSL backend?

[buffett0323]

Thanks for your interest @PKUWZP

I am wondering if Liger kernel works out of the box for Muse Glimmer?

No, the softcap form, the scale-free RMSNorms, and the Gemma-style hidden_activation field each break a different part of the generic path.

Also, this PR needs no backend-specific code and is on the same footing as Llama/Qwen w.r.t. CuTe DSL, but I'm happy to follow up with a B200 run to confirm the SM100 path is green or a PR adding the Triton fallback plus a model-level test under a non-default LIGER_KERNEL_IMPL, let me know which is useful.

[PKUWZP]

Thanks for your interest @PKUWZP

I am wondering if Liger kernel works out of the box for Muse Glimmer?

No, the softcap form, the scale-free RMSNorms, and the Gemma-style hidden_activation field each break a different part of the generic path.

Also, this PR needs no backend-specific code and is on the same footing as Llama/Qwen w.r.t. CuTe DSL, but I'm happy to follow up with a B200 run to confirm the SM100 path is green or a PR adding the Triton fallback plus a model-level test under a non-default LIGER_KERNEL_IMPL, let me know which is useful.

That makes sense. I think a B200 run would be interesting. We should merge this PR shortly after that.

[buffett0323]

Hi @PKUWZP ,
I just ran it on a B200 and found out that on SM100, the CuTe DSL path works for Muse Glimmer in bf16, and the fp32 failures are a backend-wide dtype limit, not a Muse issue.

LIGER_KERNEL_IMPL=cutedsl, B200 (SM100), same box for both columns:

	Muse Glimmer (this PR)	control: mini_llama3, mini_qwen2, mini_qwen2_vl, mini_qwen2_5_vl
bf16, through FLCE	✅ all pass	✅ 4 / 4 pass
fp32, through FLCE	❌ all fail	❌ 4 / 4 fail
fp32, not through FLCE	✅ all pass	—

Muse totals 8 passed / 6 failed, and every one of the 6 raises the same error:

RuntimeError: Native CuTe DSL FLCE requires exact SM100 hardware,
              FP16/BF16 input and weight, and mean/sum reduction.
src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py:261

Muse Glimmer needs no backend-specific code and behaves the same to Llama/Qwen on both. I'm happy to open a separate PR adding the Triton fallback for the unsupported hardware/dtype cases, plus a model-level test under a non-default LIGER_KERNEL_IMPL so this gets caught in CI.

C.C. @vaibhavjindal , how do you think about this? Thanks!!

[PKUWZP]

Hi @PKUWZP , I just ran it on a B200 and found out that on SM100, the CuTe DSL path works for Muse Glimmer in bf16, and the fp32 failures are a backend-wide dtype limit, not a Muse issue.

LIGER_KERNEL_IMPL=cutedsl, B200 (SM100), same box for both columns:

Muse Glimmer (this PR) control: mini_llama3, mini_qwen2, mini_qwen2_vl, mini_qwen2_5_vl
bf16, through FLCE ✅ all pass ✅ 4 / 4 pass
fp32, through FLCE ❌ all fail ❌ 4 / 4 fail
fp32, not through FLCE ✅ all pass —
Muse totals 8 passed / 6 failed, and every one of the 6 raises the same error:

RuntimeError: Native CuTe DSL FLCE requires exact SM100 hardware,
              FP16/BF16 input and weight, and mean/sum reduction.
src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py:261

Muse Glimmer needs no backend-specific code and behaves the same to Llama/Qwen on both. I'm happy to open a separate PR adding the Triton fallback for the unsupported hardware/dtype cases, plus a model-level test under a non-default LIGER_KERNEL_IMPL so this gets caught in CI.

C.C. @vaibhavjindal , how do you think about this? Thanks!!

This is awesome, great work! Yep, it would be great if we can have a separate PR adding the Triton fallback path. Muse Glimmer is a high profile open model, so supporting it is a big deal.