Repository navigation
fix(mtp): harden transactional cache, logprobs, and guards - #2
Conversation
b62329d to
52f54e1
Compare
AirRunner
left a comment
There was a problem hiding this comment.
Thanks for the PR!
Went through the diff, this is in good shape (MTPPromptCacheState boundary mechanism, finalize_plan bookkeeping, fail-closed guards). The verify_lp vs draft_lp fix on accept matches what we discussed.
One thing I want to note rather than just merge past: in qwen3_5.py::sanitize(), the case where mtp_num_hidden_layers > 0 but no MTP weights are present used to raise a hard ValueError. That error exists because of a report from @layer4down early on (a mlx-community quantization had mtp_num_hidden_layers: 1 in its config with no MTP tensors actually present). I see that your version replaces the hard error with a silent del self.mtp / supports_mtp = False fallback.
Actually it's not fully silent, the existing --mtp flag ignored warning in stream_generate still fires when someone actually requests --mtp on such a checkpoint, since _model_supports_mtp reads supports_mtp. So the model isn't unusable anymore, which is actually better than before (used to blocked loading the model at all, even for plain non-MTP).
However what's lost is the specific diagnosis: the old message told the user their config and weights disagree, which is what let @layer4down figure out the checkpoint was mis-packaged. The new warning just says "no MTP head," same message as a checkpoint that never had MTP support to begin with.
So I'm fine with it, but do we just want the warning message to keep the distinction or not? Curious what you think before we merge this.
|
Note: I just rebased onto upstream, so this PR's diff now shows unrelated noise on top of your change. I can't push to your branche despite maintainerCanModify=true (probably due to presence of GitHub Actions in mlx-lm or something). So if you want a clean diff you can rebase, you'll get two small conflicts in |
Rebase the selected correctness work onto AirRunner's current feat/mtp-native, preserve prompt_progress_callback wiring, and retain a specific warning when checkpoint config declares MTP but the tensors are missing. Source: bc1d114 Signed-off-by: Philip John Basile <PBasile@Basilecom.com>
476931a to
6630bde
Compare
Carry the selected correctness work from ml-explore#1740 onto AirRunner/mlx-lm:feat/mtp-native at e8ceecc.
This preserves transactional multi-turn prompt-cache alignment and finalization, target/verifier log probabilities for accepted drafts, usable-head detection, fail-closed logits-processor handling, and the focused regression coverage.
Source: PhilipJohnBasile/mlx-lm@bc1d114