[Slyne]

What does this PR do ?

In multimodal LLM training (e.g., SALM), audio/vision positions in the input sequence carry continuous features from a perception encoder that cannot be recovered by re-embedding an integer token id. The existing MTP path assumes all positions have a corresponding token id and uses embed_fn(rolled_input_ids) to produce the future-token signal — this yields embed_fn(padding_id) at every audio position, which is wrong.

This PR adds a second input mode to the MTP forward path: when pre-fused input embeddings are available (e.g., the fused audio+text tensor in SALM), the caller pre-rolls them per depth and passes them as embed_inputs, bypassing embed_fn entirely. Audio positions then carry the correct audio embedding as their future-token signal. The existing input_ids + embed_fn path is unchanged for pure-text models.

Changelog

• nemo_automodel/components/models/common/mtp/mtp.py: Updated embed_inputs docstring to document the multimodal use case alongside the existing PP use case.
• nemo_automodel/components/models/nemotron_v3/model.py: In NemotronHForCausalLM.forward(), when inputs_embeds is provided and no PP embeddings are present, pre-roll inputs_embeds per MTP depth and pass as embed_inputs, bypassing embed_fn for audio/vision positions.
• tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py: Added TestMTPInputEmbeds with three tests — embed_inputs path delivers the correct tensor to the fusion sublayer, embed_fn path is unchanged when embed_inputs is absent, and end-to-end model forward produces different mtp_per_depth_h when inputs_embeds differs from embed_tokens(input_ids).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Additional Information

This fix is required for SALM (Speech and Language Model) training in NeMo, where the input sequence interleaves text tokens and audio frame embeddings. The NeMo-side change (pre-rolling inputs_embeds before calling the LLM) will be submitted in a separate NeMo PR that depends on this one.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test d78fb9a

[akoumpa]

/ok to test fff5d7e

[akoumpa]

Thank you for the contribution @Slyne ! ,

@adil-a can you review?

Thank you

[adil-a]

/claude review

[claude]

LGTM