update lm head last token pruning - #288
Conversation
|
@copilot merge main branch into this PR. |
There was a problem hiding this comment.
Pull request overview
This PR adds an opt-in prune_lm_head flag that, when enabled, prunes the LM head computation to only the last token (prefill optimization) by inserting a Gather before the LM head. It also introduces an export_package() helper in the ORT-GenAI integration to produce a fully loadable onnxruntime-genai directory (ONNX + config artifacts) from an already-built ModelPackage.
Changes:
- Add
flags.prune_lm_head(env:MOBIUS_PRUNE_LM_HEAD) and wire it into the baseCausalLMModel.forward()path. - Add unit tests for the pruning flag behavior in the base model build path.
- Add
export_package()to the ORT-GenAI integration, refactorauto_export()to delegate to it, and add tests/docs for the new helper.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/models/base.py | Adds last-token hidden-state selection before lm_head when flags.prune_lm_head is enabled. |
| src/mobius/models/_models_test.py | Adds tests covering default vs pruned logits shape behavior for CausalLMModel. |
| src/mobius/integrations/ort_genai/auto_export.py | Adds export_package() and refactors auto_export() to use it; updates module docstring. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Adds unit tests for export_package() behavior and argument propagation. |
| src/mobius/integrations/ort_genai/init.py | Exposes export_package in the integration public API. |
| src/mobius/_model_package.py | Updates ModelPackage.save() docstring to point users to export_package() for ORT-GenAI directories. |
| src/mobius/_flags.py | Adds the prune_lm_head flag, env var support, and documentation. |
…-last-token-pruning Signed-off-by: GitHub <noreply@github.com> # Conflicts: # src/mobius/_flags.py # src/mobius/integrations/ort_genai/auto_export.py # src/mobius/integrations/ort_genai/auto_export_test.py Co-authored-by: rui-ren <15321482+rui-ren@users.noreply.github.com>
Done. Merged |
rui-ren
changed the title
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
|
The author of this PR, rui-ren, is not an activated member of this organization on Codecov. |
|
Maybe we should consider creating an option like we do for use_static_cache, as the feature flags are designed to be temporary and would eventually be removed when the feature stabilizes. Do you think it would make sense to create a new task for this? Or maybe create an option to an existing task? |
ok, sure, make sense to create a new task for this. |
|
@copilot please address comments and update from main |
Remove the MOBIUS_PRUNE_LM_HEAD feature flag and replace it with a prune_lm_head constructor parameter on CausalLMTask (and HybridCausalLMTask), mirroring the stable-API pattern already used by the static_cache option. When prune_lm_head=True, the task inserts Gather(axis=1, index=-1) followed by Unsqueeze(axis=1) on the logits output after calling module(), changing the output shape from [B, S, vocab] to [B, 1, vocab]. ONNX Runtime's graph optimizer can push this Gather backward through the LM head MatMul to avoid computing all-token logits during prefill. Update tests to use CausalLMTask(prune_lm_head=True) instead of the override_flags context manager. Also merges main branch changes (tencent_q1_0_use_native_2bit and static_cache_bias flags). Signed-off-by: Copilot <copilot@github.com> Signed-off-by: GitHub <noreply@github.com>
|
|
Done in commit Moved # Before (global flag)
with override_flags(prune_lm_head=True):
pkg = build(model_id)
# After (task option)
pkg = build(model_id, task=CausalLMTask(prune_lm_head=True))The |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
src/mobius/tasks/_causal_lm.py:332
- Same as above:
Gather(..., idx=-1)depends on negative indices. Derive the last index fromShape(logits)to avoid relying on undefined/implementation-specific behavior.
last_idx = op.Constant(value_int=-1) # scalar (rank-0) INT64
last_token_logits = op.Gather(logits, last_idx, axis=1) # [B, vocab]
logits = op.Unsqueeze(last_token_logits, op.Constant(value_ints=[1])) # [B, 1, vocab]
src/mobius/tasks/_causal_lm.py:87
- The docstring says pruning inserts a Gather "before the LM head" so only the last hidden state is projected, but the implementation below prunes after logits are computed (Gather on
[B,S,vocab]). This is misleading and also overstates the guaranteed perf win (it relies on ORT pushing the Gather backward through the projection).
prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)``
before the LM head so only the last token's hidden state is
projected to logits. Output logits shape becomes ``[B, 1,
vocab]`` instead of ``[B, S, vocab]``, reducing prefill cost
for large-vocabulary models. Set this when the downstream
src/mobius/tasks/_causal_lm.py:215
- Using
Gather(..., idx=-1)relies on negative indices, which are not guaranteed by the ONNX Gather spec across runtimes. Compute the last index fromShape(logits)instead to keep the graph spec-compliant and portable.
This issue also appears on line 330 of the same file.
last_idx = op.Constant(value_int=-1) # scalar (rank-0) INT64
last_token_logits = op.Gather(logits, last_idx, axis=1) # [B, vocab]
logits = op.Unsqueeze(last_token_logits, op.Constant(value_ints=[1])) # [B, 1, vocab]
src/mobius/tasks/_causal_lm.py:99
- The PR description says this is enabled via
MOBIUS_PRUNE_LM_HEAD=1orflags.prune_lm_head=True, but the implementation only supports passingprune_lm_headwhen constructingCausalLMTask/HybridCausalLMTask. If the env-var / global flags integration is a requirement, it’s currently missing.
def __init__(
self,
*,
static_cache: bool = False,
max_seq_len: int | None = None,
prune_lm_head: bool = False,
):
self._static_cache = static_cache
self._max_seq_len = max_seq_len
self._prune_lm_head = prune_lm_head
src/mobius/models/_models_test.py:233
test_prune_emits_gather_on_logitsonly checks the inferred output shape. That can miss regressions where shape inference still yields[B,1,V]but the pruning ops are not actually wired (or get optimized away unexpectedly). Making the test assert theUnsqueeze <- Gatherproducer chain keeps it directly tied to the intended graph edit.
logits = next(v for v in model.graph.outputs if v.name == "logits")
# Logits must still be rank-3 [B, 1, vocab] (NOT rank-4 [B, 1, 1, V])
assert len(logits.shape) == 3, (
f"Expected rank-3 logits [B, 1, V] after pruning, got rank "
f"{len(logits.shape)}: shape={list(logits.shape)!r}"
)
# Pruned: dim 1 must be the literal integer 1
seq_dim = logits.shape[1]
assert seq_dim == 1, f"Expected logits dim 1 to be 1 after pruning, got {seq_dim!r}"
# Last dim is still the vocabulary size
config = make_config()
assert logits.shape[2] == config.vocab_size
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| assert logits.shape[2] == config.vocab_size | ||
|
|
||
| def test_prune_does_not_change_input_shapes(self): | ||
| """Pruning only affects output; input_ids still has dynamic sequence_length |
PR Description
Adds an opt-in
prune_lm_headflag to Mobius, mirroring ONNX Runtime GenAI Model Builder behavior.When enabled (
MOBIUS_PRUNE_LM_HEAD=1orflags.prune_lm_head = True), Mobius inserts aGather(axis=1, idx=-1)before the LM head so only the final token logits are computed during prefill.This reduces unnecessary
[B, S, vocab]logit computation and improves prefill performance.Default is
Falseto preserve compatibility with:Currently supported by models using the base
CausalLMModel.forward()path (Llama, Mistral, Qwen2/2.5/3, etc.). Models with customforward()implementations currently ignore the flag.