Fix MTP cache overflow during generation

#19
by iCavan - opened

The auxiliary MTP layer is used for training, but it was also being executed during normal inference. With use_cache=True, it uses layer index 42 after the 42 inference layers have populated cache slots 0–41, so standard generate() fails with an IndexError.

This change runs MTP only when labels are provided, keeps it out of the inference cache, and preserves MTP logits and loss for training. It also avoids mutating the MTP hidden-state list while iterating over it. The main LM logits are unchanged because they are computed from main_hidden_states before the auxiliary MTP layer.

Validation:

  • python -m py_compile modeling_bailing_moe_v3.py
  • Tiny-model regression covering cached generation, inference cache length, and the MTP training path
  • Full Ling 3.0 Flash checkpoint on 4Γ— H20-3e in BF16
  • All 43 modules remained loaded; a fail-fast hook confirmed the MTP layer was not called during inference
  • Standard cached greedy generation completed for 32 tokens in 23.09 seconds
  • Generated tokens matched the pre-fix no-MTP semantic baseline 32/32

Skipping the training-only layer also removes its unnecessary inference work instead of extending the cache to accommodate a head whose output is not used for generation.

zheyishine changed pull request status to merged

Sign up or log in to comment