Potential Issue with Qwen2VisionModel RoPE

[xiaoxiao26]

I ran into an issue today with the RoPE embeddings for Qwen2VisionModel

When the RoPE frequency embeddings are constructed here: https://github.com/NVIDIA-NeMo/NeMo/blob/main/nemo/collections/vlm/qwen2vl/model/vision.py#L416-L444, it is of length total number of tokens T in the batch of images

This can lead to issues in the global attention layers for mixed resolution batches in downstream Megatron repo. For example, lets say we have a batch where first image is 640 tokens and second image is 960 tokens. Then when indexing into the RoPE embeddings for the second image, it will index the first 960 tokens rather than the last 960 tokens in the frequency tensor, which is a mismatch: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L140

The local attention layers are fine because the patch lengths are the same, but this makes global attention layers wrong whenever you have a batch of images of mixed resolutions. This issue seems to occur for both the fused and unfused RoPE implementations Megatron.

[gtangg12]

Also hit same issue!

[yaoyu-33]

Thanks for reporting, confirm it's a bug. Will try to file a fix.

[yaoyu-33]

Hi @xiaoxiao26 @gtangg12 : I am checking the issue right now.
I notice line: https://github.com/NVIDIA-NeMo/NeMo/blob/main/nemo/collections/vlm/qwen2vl/model/vision.py#L497 has shape torch.Size([S, 1, 1280]) because we flattened everything with seq packing. where dim 0 of x should cover the entire range of position ids. Can you share a bit more about your reproducer.

[gtangg12]

Hey, in Megatron we eventually reach this:

https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L238
and
https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L143
https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L140
which assumes that freqs is 1D rope instead of MMRope and that freqs is the longest packed sequence

The python code branch can be fixed/optimized by using _apply_rotary_pos_emb_bshd for cp_size = 1 and some vectorization for cp_size > 1 (the for loop and torch.split is needed for the max_seq definition for freqs but if its aligned as nemo assumes, we don't need it and can directly do the aforementioned). However, the cuda kernel for fused rope we have no idea how to fix/digs into TE. Can you please take a look?

[zhenyih]

ai generated, please verify

The issue in Qwen2VisionModel occurs when processing batches with mixed resolution images. Currently, when RoPE frequency embeddings are constructed, they're created as a single tensor for all tokens in the batch. For mixed resolution batches (like a 640-token image followed by a 960-token image), Megatron's global attention layers incorrectly index into these embeddings.

The problem happens because when indexing the second image's embeddings, the code uses the first 960 tokens of the frequency tensor instead of starting at position 640. This creates a position mismatch in global attention layers, while local attention layers work fine since they operate within individual images.

To fix this, the rot_pos_emb method needs to be modified to ensure position IDs are properly offset for each image in the batch, allowing Megatron's global attention layers to correctly index the RoPE embeddings.

[yaoyu-33]

@gtangg12 @zhenyih I try to understand where the core issue is, is it from the full attn vs. local attn blocks?

our colleague also has another implementation of megatron version vl at https://github.com/ISEEKYAN/mbridge/blob/main/mbridge/models/qwen2_5_vl/vision_model.py it handles the packed seq params slightly differently maybe due to the same issue.

for cp=1 case, it will return freqs[: x.size(0)] from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L140. This length is correct right, because x here is of dimension [S, 1, 1280] and S is the total length of all images. there shouldn't be any cutoff?

Beyond this issue, just FYI we are working on a new approach:
by mostly re-using HF code for vision part of model, as vision model is relatively small and doesnt need complex model parallel. NVIDIA-NeMo/Megatron-Bridge#558

in you use cases, is it reasonable to have vision part go with torch backend instead of megatron? That will overcome many implementation confusion.

[xiaoxiao26]

@yaoyu-33 regarding

"for cp=1 case, it will return freqs[: x.size(0)] from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L140. This length is correct right, because x here is of dimension [S, 1, 1280] and S is the total length of all images. there shouldn't be any cutoff?"

My understanding is that x here is actually one image, not all images, because it is split to individual images here: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/embeddings/rope_utils.py#L181

[yaoyu-33]

@xiaoxiao26 : thanks for your comments. Seems the fix needs to be ported to megatron core side.

[gtangg12]

@yaoyu-33 it seems that this issue is also mirrored for the CUDA kernel fused_apply_rotary_pos_emb_thd. Can you please take a look?

[yaoyu-33]

I am adding a fix to mcore internally, replacing the logic of _get_thd_freqs_on_this_cp_rank. Wonder if the fix below can help to solve your issues?

def _get_thd_freqs_on_this_cp_rank(cp_rank: int, cp_size: int, x: Tensor, freqs: Tensor, offset: int = 0) -> Tensor:
    """Get the correct frequency slice for this context parallel rank with optional sequence offset.
    
    Args:
        cp_rank: Current context parallel rank
        cp_size: Total context parallel size  
        x: Input tensor for current sequence
        freqs: Frequency tensor - either full batch positions or max sequence length
        offset: Starting position offset for this sequence in the original batch (default: 0)
        
    Returns:
        Tensor: Frequency slice corresponding to this CP rank's portion of the sequence
        
    Note:
        This function supports two modes based on the offset parameter:
        1. offset > 0: Exact mapping mode - freqs contains all positions across all sequences.
           The offset ensures each sequence gets frequencies from its actual position within
           the overall batch. Critical for non-1D RoPE in VLMs where spatial positions matter.
        2. offset = 0: Traditional mode - freqs contains only max sequence length positions.
           All sequences use frequencies starting from position 0, preserving backward compatibility.
    """
    if cp_size > 1:
        cp_seg = x.size(0) // 2
        full_seqlen = cp_size * x.size(0)
        # Apply offset to both forward and backward segments for context parallelism
        # offset=0: traditional behavior, freqs[0:cp_seg] and freqs[...] 
        # offset>0: exact mapping, freqs[offset+0:offset+cp_seg] and freqs[offset+...]
        return torch.cat(
            [
                freqs[offset + cp_rank * cp_seg : offset + (cp_rank + 1) * cp_seg],
                freqs[offset + full_seqlen - (cp_rank + 1) * cp_seg : offset + full_seqlen - cp_rank * cp_seg],
            ]
        )
    else:
        # For single context parallel rank: 
        # offset=0: use freqs[0:x.size(0)] (traditional)
        # offset>0: use freqs[offset:offset+x.size(0)] (exact mapping)
        return freqs[offset : offset + x.size(0)]


def _apply_rotary_pos_emb_thd(
    t: Tensor,
    cu_seqlens: Tensor,
    freqs: Tensor,
    rotary_interleaved: bool = False,
    multi_latent_attention: bool = False,
    mscale: float = 1.0,
    cp_group: torch.distributed.ProcessGroup = None,
) -> Tensor:
    """A baseline implementation of applying RoPE for `thd` format.

    Args:
        t (Tensor): Input tensor T is of shape [t, h, d]
        cu_seqlens(Tensor):  Cumulative sum of sequence lengths in a batch for `t`,
        with shape [b + 1] and dtype torch.int32.
        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d]
        cp_group (torch.distributed.ProcessGroup): The context parallel group

    Returns:
        Tensor: Shape [t, h, d]. The input tensor after applying RoPE.
    """

    if cp_group is None:
        raise ValueError("cp_group must be provided for THD format RoPE")
    cp_size = cp_group.size()
    cp_rank = cp_group.rank()
    seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist()

    # Handle two different frequency tensor formats:
    # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences
    #    -> Use offset-based mapping for exact positional correspondence
    # 2. Otherwise: freqs contains only max sequence length positions  
    #    -> Use traditional mapping without offsets (map first :seqlen part)
    if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]:
        # CASE 1: Exact mapping with offsets
        # freqs tensor contains positional embeddings for all positions across all sequences
        # Each sequence gets frequencies from its actual position within the overall batch
        sequence_splits = torch.split(t, seqlens)
        results = []
        for i, x in enumerate(sequence_splits):
            # Calculate the starting position of this sequence in the original batch
            # cu_seqlens[i] gives us the cumulative offset where this sequence begins
            seq_start_offset = cu_seqlens[i].item()
            
            results.append(
                _apply_rotary_pos_emb_bshd(
                    x.unsqueeze(1),
                    # Pass the sequence-specific offset to get the correct frequency slice
                    _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset),
                    rotary_interleaved=rotary_interleaved,
                    multi_latent_attention=multi_latent_attention,
                    mscale=mscale,
                )
            )
        
        return torch.cat(results).squeeze(1)
    else:
        # CASE 2: Traditional mapping without offsets
        # freqs tensor contains only max sequence length, all sequences use freqs[0:seqlen]
        # This is the original behavior for backward compatibility
        return torch.cat(
            [
                _apply_rotary_pos_emb_bshd(
                    x.unsqueeze(1),
                    _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs),
                    rotary_interleaved=rotary_interleaved,
                    multi_latent_attention=multi_latent_attention,
                    mscale=mscale,
                )
                for x in torch.split(t, seqlens)
            ]
        ).squeeze(1)

[yaoyu-33]

@gtangg12 : from the description, I think the issue comes from the incorrect freqs are paired with the seqs.
with the correction above, do you see still issues?

if the fused version fused_apply_rotary_pos_emb_thd has different behavior than apply_rotary_pos_emb_thd we can start an issue in Transformer Engine to track that.

[yaoyu-33]

@xiaoxiao26 @gtangg12 wonder if you have any comments on this potential fix?

[gtangg12]

I think your implementation is fine, but can be optimized. The main blocker is the for loop. For the cp=1 case I think you remove the loop like this:

return _apply_rotary_pos_emb_bshd(
    t.unsqueeze(1),
    freqs,
    rotary_interleaved=rotary_interleaved,
    multi_latent_attention=multi_latent_attention,
    mscale=mscale,
).squeeze(1)

For cp > 1 case you can write some vectorization code to remove the loop.

Doing this increases the training speed 5x, brining it closer to the fused performance. I believe the fused kernel has the same issue, and we should check that.