Skip to content

Fix Phi3/Phimoe prepare_inputs_for_generation forwarding logits_to_keep=None#47531

Open
pwasiewi wants to merge 1 commit into
huggingface:mainfrom
pwasiewi:fix-phi3-phimoe-logits-to-keep-none
Open

Fix Phi3/Phimoe prepare_inputs_for_generation forwarding logits_to_keep=None#47531
pwasiewi wants to merge 1 commit into
huggingface:mainfrom
pwasiewi:fix-phi3-phimoe-logits-to-keep-none

Conversation

@pwasiewi

@pwasiewi pwasiewi commented Jul 24, 2026

Copy link
Copy Markdown

CI

What does this PR do?

Fixes a generation crash in Phi3 and Phimoe: their prepare_inputs_for_generation overrides default logits_to_keep=None and forward it unconditionally into forward(). When GenerationMixin._supports_logits_to_keep() cannot introspect a wrapped forward() (e.g. TRL/PEFT training stacks wrap forward without functools.wraps), generate() never injects its default of 1, so None reaches forward() — where None as a slice index turns the lm_head slice into an unsqueeze (hidden_states[:, None, :]), producing 4-dim logits and crashing sampling with RuntimeError: prob_dist must be 1 or 2 dim in torch.multinomial.

The fix coerces None -> 1 in the two overrides (modular + regenerated modeling files): generation only consumes the last position, and assisted decoding overrides model_inputs["logits_to_keep"] afterwards, so it is unaffected. Phi4Multimodal already avoids the bug by defaulting to 0 in its own override; all non-overriding models are unaffected because the kwarg is simply absent.

Fixes #47530

Before submitting

Who can review?

@gante

…ep=None

Phi3ForCausalLM and PhimoeForCausalLM override prepare_inputs_for_generation
with a logits_to_keep=None default and forward it unconditionally. When
GenerationMixin._supports_logits_to_keep() cannot introspect a wrapped
forward() (e.g. TRL/PEFT training stacks wrap forward without
functools.wraps), generate() never injects its default of 1, so None reaches
forward(). There, None as a slice index turns the lm_head slice into an
unsqueeze (hidden_states[:, None, :]) producing 4-dim logits, which crashes
sampling with "RuntimeError: prob_dist must be 1 or 2 dim" in
torch.multinomial.

Coerce None -> 1: generation only consumes the last position, and assisted
decoding overrides model_inputs["logits_to_keep"] afterwards, so it is
unaffected. Phi4Multimodal already avoids this by defaulting to 0.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: phi3, phimoe

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 30133295822
Result: success | Grafana metrics are not available yet.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Phi3/Phimoe: prepare_inputs_for_generation forwards logits_to_keep=None into forward() -> 4-dim logits -> torch.multinomial crash

1 participant