[bohnstingl]

Overview

Replace the model-local MuseGlimmerRMSNorm (added in #51655) with the existing RMSNorm and GemmaRMSNorm CustomOps. The output is bit-identical in bf16.

Claims

• Removes the model-local MuseGlimmerRMSNorm class. All norms in the model now use the shared CustomOps, so they get are enabled for the OOT-platform swapping.
• No numerical change: each replacement matches the old class bit-for-bit in bf16.
• No checkpoint or weight-loading change: parameter names and shapes are unchanged.

Validation

Old class vs replacement, bf16, A100, x = 257×5120, both forward_native and forward_cuda:

Norm (old config)	Replacement	Result
with_scale=False (qk_norm, embed_norm, perception_emb_norm)	RMSNorm(dim, has_weight=False)	bit-exact
weight_offset=1 (4 sandwich norms)	GemmaRMSNorm(dim)	bit-exact
final norm	RMSNorm(dim, dtype=torch.float32)	bit-exact

.venv/bin/python -m pytest tests/models/test_initialization.py \
    -k "MuseGlimmerForCausalLM or MuseGlimmerForConditionalGeneration" -v
# 2 passed

gsm8k on meta-models/Muse-Glimmer-30B (BF16, TP=2 on 2×A100-80GB, --model-impl vllm, 5-shot, 1319 questions, greedy, max 256 tokens):

	Accuracy	Correct	Invalid	Throughput (tok/s)
main (e318ac0979)	0.9181	1211 / 1319	1	1911
this PR	0.9189	1212 / 1319	0	1930

The 1-question delta is batching noise; the norms themselves are bit-exact.

vllm serve meta-models/Muse-Glimmer-30B --model-impl vllm --dtype bfloat16 \
    --tensor-parallel-size 2 --max-model-len 4096 \
    --limit-mm-per-prompt '{"image": 0, "video": 0}'
python tests/evals/gsm8k/gsm8k_eval.py --num-shots 5 --num-questions 1319

Details

The old class covered three cases, each of which already exists in-tree:

• No weight: norm(x.float()).type_as(x) is RMSNorm(has_weight=False).
• (1 + w) sandwich norms: these are HF's MuseGlimmerTextCenteredRMSNorm, a copy of Gemma's norm, so GemmaRMSNorm matches.
• Final norm: HF computes (norm(x.float()) * w.float()).type_as(x). The in-tree op multiplies in the weight's dtype, so the weight is kept in fp32 to match HF's single cast. The fp32 weight survives the bf16 default dtype at init and default_weight_loader. The trade-off is that the fused vllm_c kernel requires matching weight/input dtypes, so this one norm (once per forward) takes the native PyTorch path instead.

Not in scope: Gemma2-style fused residual-add in the decoder layer, which would change forward.

No open PR touches this: searched open PRs for MuseGlimmerRMSNorm, MuseGlimmer RMSNorm and muse_glimmer norm.

cc @xianbaoqian @jeejeelee @zixi-qi @TheEpicDolphin (authors of #51655)

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[jeejeelee]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99135 for commit f4fe8c8ff395.

ℹ️ This PR is 48 commits behind upstream main. Merged upstream main into this branch as f4fe8c8ff395 before starting CI. Pull the branch before pushing more changes.

[bohnstingl]

I think the failing CI tests are not from this PR and a resolution is provided through #65140

[bohnstingl]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99205 for commit 59c76771e846.

[jeejeelee]

@bohnstingl if so, we can force merge this PR, please don't trigger CI again