Mistral3ForConditionalGeneration: forward does not declare logits_to_keep, so fused losses are silently replaced

Describe the bug

Mistral3ForConditionalGeneration.forward does not declare logits_to_keep; it
absorbs it into **kwargs. This is the same defect as #3510 (MiniMax-M3), in a
different registered model.

nemo_automodel/components/models/mistral4/model.py L833:

def forward(
    self,
    input_ids: torch.Tensor | None = None,
    *,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    image_sizes: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    **kwargs: Any,
) -> torch.Tensor:

Because _supports_logits_to_keep inspects the forward signature
(components/utils/model_utils.py L76-77), the check returns False and the
recipe silently swaps the configured loss:

if not _supports_logits_to_keep(probe_module) and not isinstance(loss_fn, MaskedCrossEntropy):
    logger.warning("logits_to_keep not found in model.forward. Using MaskedCrossEntropy instead.")
    return MaskedCrossEntropy()

(recipes/llm/train_ft.py L159, recipes/vlm/finetune.py L554, and
_transformers/infrastructure.py L640.)

The class is registered, so this is user-reachable —
_transformers/registry.py L194-197:

(
    "Mistral3ForConditionalGeneration",
    ("nemo_automodel.components.models.mistral4.model", "Mistral3ForConditionalGeneration"),
),

The kwarg is not honoured anywhere in the class either: logits_to_keep is
never referenced in the class body, and forward ends with

logits = lm(hidden_states) if lm is not None else hidden_states

so the head is applied over the full sequence regardless.

The two consequences from #3510 apply unchanged:

1. The full logits tensor is materialised. Mistral4Config.vocab_size
   defaults to 131072 (mistral4/configuration.py L30), so the head output is
   [tokens, 131072].
2. The configured loss_fn's settings are dropped with it. The replacement
   is a fresh MaskedCrossEntropy(), so a recipe that sets
   fp32_upcast: false loses it and the same tensor is upcast to fp32.

Why this looks like an oversight rather than a design choice

• Its own text backbone, Mistral4ForCausalLM, declares logits_to_keep — same
  file, L417.
• Its sibling in the same registry block,
  Mistral3FP8VLMForConditionalGeneration, declares it
  (mistral3_vlm/model.py L206).
• I scanned every *ForCausalLM / *ForConditionalGeneration /
  *ForImageTextToText class under components/models (55 classes). All but
  these declare it explicitly, and the classes that inherit forward from a
  Hugging Face parent (Gemma4UnifiedForConditionalGeneration,
  Gemma4DrafterForCausalLM) inherit a signature that declares it, so they are
  fine.

Steps/Code to reproduce bug

Configure any Mistral 4 multimodal recipe with a fused loss:

model:
  _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained
  pretrained_model_name_or_path: <mistral3-vlm-checkpoint>

loss_fn:
  _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy

Training logs logits_to_keep not found in model.forward. Using MaskedCrossEntropy instead. and the configured loss is discarded.

Signature check without a checkpoint:

import inspect
from nemo_automodel.components.models.mistral4.model import Mistral3ForConditionalGeneration

print("logits_to_keep" in inspect.signature(Mistral3ForConditionalGeneration.forward).parameters)
# False

Expected behavior

Mistral3ForConditionalGeneration.forward declares logits_to_keep and honours
it — slicing hidden_states before the head — so a configured
FusedLinearCrossEntropy is kept instead of being silently replaced.

Environment overview

• main at 3ddef9b1. Static signature/registry inspection on CPU; no GPU or
  checkpoint needed to observe the False.

Additional context

One more to confirm, lower confidence: MiniMaxM3SparseForCausalLM
(minimax_m3_vl/model.py L322) has the same gap. It is the sibling of the class
#3510 fixed, in the same file — MiniMaxM3SparseForConditionalGeneration now
declares it, this one does not. It is not in MODEL_ARCH_MAPPING, so it may be
intentionally an internal text backbone, but it is exported from
minimax_m3_vl/__init__.py and is the value of architectures in
minimax_m3_vl/config.py L48. Worth a maintainer's eye on whether it is
reachable as a top-level model.

Happy to send a PR for the Mistral case (declare and honour the kwarg, plus a
signature test), and to include the MiniMax one if you confirm it should be
top-level.

Start in nemo_automodel/components/models/mistral4/model.py around Mistral3ForConditionalGeneration.forward at L833, comparing it with Mistral4ForCausalLM around L417. Then read _supports_logits_to_keep in components/utils/model_utils.py and the loss selection in recipes/llm/train_ft.py and recipes/vlm/finetune.py. Done means the signature check succeeds, logits_to_keep is honored before the head, and the configured fused loss is no longer replaced.