mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Buffett LiuandZhipeng Wang
e6a81bb0c3
## 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: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Co-authored-by: Zhipeng Wang <[email protected]>