Skip to content

MTP Multimodal compatible #8303

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jhvmhg wants to merge 3 commits into from
Closed

MTP Multimodal compatible #8303

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c38b406
MTP Multimodal compatible
jhvmhg Mar 12, 2026
a8e2a4c
Apply suggestion from @gemini-code-assist[bot]
jhvmhg Mar 12, 2026
fb9a8db
MTP Multimodal compatible
jhvmhg Mar 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions swift/megatron/init.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,7 +430,6 @@ def forward(
Union[Tensor, Tuple[Tensor, Tensor]]: The output hidden states tensor of shape
[s, b, h], and optionally the updated context tensor if cross-attention is used.
"""
# TODO: Multimodal compatible
assert context is None, 'multi token prediction + cross attention is not yet supported.'
input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings(
input_ids=input_ids,
Expand All @@ -441,9 +440,15 @@ def forward(
packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'
if packed_seq:
assert not self.transformer_layer.self_attention.config.apply_rope_fusion
assert position_ids.shape[0] == 1, f'position_ids.shape: {position_ids.shape}'
rotary_pos_emb = rotary_pos_emb[position_ids[0]]
if self.config.recompute_granularity == 'full' and self.training:
if position_ids.shape[0] == 1:
# Standard RoPE: index the rope table by shifted position_ids
rotary_pos_emb = rotary_pos_emb[position_ids[0]]
elif rotary_pos_emb is not None:
# MRoPE: rotary_pos_emb is already per-token [seq, bs, 1, dim],
# roll along sequence dim to get embeddings for shifted positions
rotary_pos_emb = torch.roll(rotary_pos_emb, shifts=-1, dims=0)
_do_recompute = self.config.recompute_granularity == 'full' and self.training
if _do_recompute:
hidden_states = self._checkpointed_forward(
partial(
self._proj_and_transformer_layer,
Expand Down