fix(minimax_m3_vl): declare and honour logits_to_keep so fused losses are kept (#3511)
* fix(minimax_m3_vl): declare and honour logits_to_keep so fused losses are kept
MiniMaxM3SparseForConditionalGeneration.forward absorbed logits_to_keep into
**kwargs, so _supports_logits_to_keep (which inspects the signature) returned
False and the recipe silently replaced any memory-efficient loss with a fresh
MaskedCrossEntropy():
if not _supports_logits_to_keep(model) and not isinstance(self.loss_fn, MaskedCrossEntropy):
logger.warning("logits_to_keep not found in model.forward. Using MaskedCrossEntropy instead.")
self.loss_fn = MaskedCrossEntropy()
Two consequences: the [tokens, vocab_size] logits tensor is materialised (200064
vocab, so tens of GB at long context), and the configured loss_fn's own settings
-- including fp32_upcast: false -- are dropped with it, upcasting that tensor to
fp32. Callers already pass logits_to_keep=1 unconditionally, so the parameter
was accepted and then ignored.
The signature check is deliberately stricter than _supports_seq_lens (which does
accept **kwargs): absorbing the kwarg does not imply the model returns hidden
states, so an explicit parameter is the right contract. Every other custom model
in components/models declares it -- minimax_m2, glm_moe_dsa, deepseek_v4,
gemma4_moe, qwen3_moe, gpt_oss, llama, qwen3_next, glm4_moe, step3p7 -- M3 was
the only one missing it.
Requesting the fused path together with MTP raises instead of returning a dict
mtp_logits cannot consume, matching the MTP-under-PP guard earlier in the same
forward.
MiniMaxM3SparseForCausalLM has the same gap but is not in the architecture
registry, and its thd branch reshapes the return value, so it is left alone here.
Signed-off-by: Xinyu Li <xli74@kent.edu>
* test(minimax_m3_vl): cover the logits_to_keep contract
Five CPU tests over the existing tiny-config fixtures:
- the recipe's own _supports_logits_to_keep probe returns True, which is the
check that silently downgraded the loss when the kwarg was absorbed by
**kwargs, so this is the regression that matters most;
- logits_to_keep returns a mapping with hidden_states shaped [.., hidden_size]
rather than [.., vocab_size], i.e. lm_head was skipped;
- omitting it still returns the logits tensor, so the default path is unchanged;
- MTP plus logits_to_keep raises NotImplementedError;
- MTP without logits_to_keep still produces logits and mtp_per_depth_logits.
Adds an mtp_vlm_model fixture: the shared mtp_model builds
MiniMaxM3SparseForCausalLM, while the branch under test lives on the
conditional-generation class.
Signed-off-by: Xinyu Li <xli74@kent.edu>
* Update return type for MiniMaxM3CausalLMOutput.forward
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
* style(minimax_m3_vl): indent the forward return annotation for ruff format
ruff format wants the closing paren of the multi-line signature at the method's
indentation. Valid Python either way, but `ruff format --check` fails on it.
Only this file is touched: upstream/main already has 12 files that ruff 0.16.2
would reformat, so those are left alone rather than folded into this PR.
Signed-off-by: Xinyu Li <xli74@kent.edu>
---------
Signed-off-by: Xinyu Li <xli74@kent.edu>
Co-authored-by: Abhishree Thittenamane <47577437+athitten@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> C
coderaBruce committed
e7b3451d653042f216b16152fc23912c1d2563c7
Parent: 2afe8b7
Committed by GitHub <noreply@github.com>
on 8/14/2026, 6:51:00 PM