[ThoreKoritzius]

Overview

Supports ModernBertDecoderForCausalLM (the Ettin decoders, 17M–1B) through the Transformers modeling backend. To do this, the backend's causal LM head learns to handle output heads that transform the hidden states before the output embeddings.

There is no existing issue or PR for this. I searched open and closed issues and PRs for ModernBertDecoder, ettin, modernbert decoder, and for output-head / lm_head changes to the Transformers backend.

Claims

• jhu-clsp/ettin-decoder-* models load and generate with vLLM. Before this change, loading failed with ValueError: There is no module or parameter named 'lm_head.dense' in TransformersForCausalLM.
• Generic: the Transformers backend now supports any ForCausalLM whose output embeddings are not called lm_head, plus one module applied to the hidden states before them. For example, decoder(lm_head(hidden_states)) with a biased, tied decoder.
• No cost or behavior change for other models. Only classes that override get_output_embeddings are inspected: 31 of the 178 AutoModelForCausalLM classes in Transformers 5.19. The other 147 return early without building anything. The inspected ones that already work keep their behavior; for example, GPT-NeoX's test_common still passes.

Validation

H100 80GB, VLLM_USE_PRECOMPILED=1 editable install on top of ffec221276, Transformers 5.19.0, torch 2.13.

Correctness against Transformers. Greedy, 32 tokens, the 8 prompts in tests/prompts/example.txt, float32 on both sides:

model	identical generations
jhu-clsp/ettin-decoder-17m	8 / 8 (32/32 tokens each)
jhu-clsp/ettin-decoder-150m	7 / 8 (the 8th diverges only at its last token, 32/32)

In bfloat16, check_logprobs_close against HfRunner fails for both sizes. At each mismatch the two candidate tokens have identical vLLM logprobs (e.g. detect / identify both −3.1181, performance / throughput both −0.8469). That is a bf16 near-tie in small models, not a model difference, so the end-to-end test runs in float32.

Tests (all passing):

# new: head detection + renaming (meta device, no GPU), and HF-vs-vLLM float32 e2e on ettin-decoder-17m
pytest tests/models/transformers/test_backend.py -k test_output_head          # 4 passed
# new registry entry initializes with dummy weights
pytest tests/models/test_initialization.py -k ModernBertDecoder               # 1 passed
# regression: an architecture that overrides get_output_embeddings and already runs on this backend
pytest tests/models/language/generation/test_common.py -k "pythia and not prompt_embeds_True"   # 2 passed
pre-commit run --files <changed files>                                        # all hooks pass (incl. mypy)

Details

Root cause. The backend builds AutoModel (the base model) and adds its own ParallelLMHead named lm_head. It assumes the Transformers ForCausalLM is lm_head(model(...)). ModernBERT decoder is instead decoder(lm_head(model(...))):

• its lm_head is a prediction head (dense → GELU → LayerNorm);
• decoder is the output projection, tied to the input embeddings and with its own bias.

So the checkpoint's lm_head.dense.* had nowhere to load, and skipping the prediction head would also produce wrong logits.

Change (transformers/causal.py, about 90 lines):

• CausalMixin._output_head finds the output head:
  • It returns at once unless the ForCausalLM class overrides get_output_embeddings; only such a class can name its output embeddings differently.
  • Otherwise, it builds that class on meta, as SequenceClassificationMixin already does for classifiers. It takes the output embeddings' attribute name, plus the one parameterized child that is neither the base model nor the output embeddings, if there is exactly one.
  • With more than one such child, the order of application can't be known. The backend logs a warning and keeps today's behavior.
• load_weights renames the output embeddings' weights to lm_head.* and the transform's weights to lm_head_transform.*. The existing lm_head.bias handling then loads decoder.bias unchanged.
• compute_logits applies lm_head_transform, if present, before the logits processor. It runs only on the sampled positions and outside the compiled region. It adds no CPU/GPU synchronization.

The transform stays a plain Transformers module, replicated on every TP rank. It is a hidden×hidden matrix, so this is negligible.

Why the Transformers backend and not a native model file. Small dense models are being moved to the backend (#48100, #53615, #59679, #59701). This architecture runs on it once the head is understood: separate q/k/v projections are fused into QKVParallelLinear, and per-layer sliding windows are mapped from layer_types. A native implementation would be about 300 lines for one architecture. For reference, I wrote one first (out of tree, as a plugin), and it matches this PR's outputs.

Who else this helps. These architectures also overrode get_output_embeddings with a non-lm_head name or a transform head in the scan:

• BioGpt (output_projection)
• Electra (generator_predictions + generator_lm_head)
• GPTNeoXJapanese (embed_out)
• TrOCR (output_projection)

They are not registered here and I did not test them.

Limitations and untested paths:

• Single GPU only. I had no multi-GPU machine.
• TP > 1: ModernBERT decoder has no TP plan, so it relies on the backend's fallback: QKV and attention are sharded, everything else replicated. Wi holds the GLU input and gate halves in one matrix; it is not fused, so it stays replicated, which is correct but not optimal.
• PP: the transform is only built on the last rank, next to lm_head.
• Docs: the LoRA and PP checkmarks in supported_models.md follow the backend's general support, like the other backend entries. They were not exercised for this model.
• Fusion count: the GLU fusion count in the e2e test is (0, 1) because Wi is already a single matrix.

AI assistance. This PR was written with AI assistance (Claude). I reviewed every changed line and ran the tests above myself.

---

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.

🤖 Generated with Claude Code

https://claude.ai/code/session_01G6Bm8s1uyizXtz9t5ZCZAw

[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.

[mergify]

Documentation preview: https://vllm--61246.org.readthedocs.build/en/61246/

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[hmellor]

Thanks for this PR!

We can definitely get these models in, but I think the change is going to be a bit more involved because this PR currently assumes a few things that aren't universal across models. I'll make a PR handling the detection of extra things in the causal model

[ThoreKoritzius]

Thanks for this PR!

We can definitely get these models in, but I think the change is going to be a bit more involved because this PR currently assumes a few things that aren't universal across models. I'll make a PR handling the detection of extra things in the causal model

Makes sense to generalise this, thanks for having a look into that! Would then rebase on #61276 to adapt the registry, will test in the meantime, happy to help with anything else too.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ThoreKoritzius.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork