RoPE positions are never set

[francesco-bertolotti]

I am looking into how RoPE positions are handled. I noticed that the models support a positions field:

torchtitan/torchtitan/models/qwen3/model.py

Lines 57 to 63 in 5732118

	def forward(
	self,
	x: torch.Tensor,
	freqs_cis: torch.Tensor,
	attention_masks: AttentionMasksType | None,
	positions: torch.Tensor | None = None,
	):

However, when examining the dataloader and dataset implementation, it appears that positions is always None.

If document packing is used, this could lead to misaligned RoPE embeddings across documents. Specifically, the first token of a packed document might not receive the first positional embedding, because the position index continues from the previous document in the packed sequence. As a result, the RoPE embedding applied to the first token of a document may not correspond to position 0 of that document.

Is this expected behavior, or should positions be reset per document when packing is enabled?

[tianyu-l]

It sounds a bug to me. Whenever we use document packing, the dataloader should also yield positions to be consumed by RoPE. Does that sound right?

cc @joecummings who's doing related work.

[francesco-bertolotti]

Makes sense. I was curious if that was a design choice, since RoPE is all about those relative position encodings. Maybe the exact positioning wasn't a huge factor here?

[joecummings]

This is definitely a bug, albeit likely a fairly small one like you pointed out b/c the relative position between two tokens within a document will be the same. Looks like position ids weren't even introduced until #2112. I can try to get something up for this later today or tomorrow unless you (@francesco-bertolotti) want to take a crack at it?

[francesco-bertolotti]

I think I can work at it this weekend. Let me know if it is not too late for you

[shuhuayu]

For the case of sample packing, do we also need to use the flex attention backend with block_causal to prevent cross sample attention.

[tianyu-l]

today torchtitan:

• txt dataloading only has packing mode
• flex / varlen supports document packing; sdpa doesn't

From correctness perspective, combining txt dataloading with sdpa is NOT correct. But it may be still nice to showcase it in torchtitan. @joecummings maybe we can introduce a pack mode bool in dataloading and error out if packing + sdpa are used together.

[francesco-bertolotti]

Quick question for @joecummings and @tianyu-l regarding RoPE when using sample packing.

If a sample spans multiple chunks (e.g., the sequence is truncated at the end of one chunk and continues in the next), should the RoPE positions reset to zero for the new chunk, or should they continue from the last position before truncation?

My understanding is that if the sequence logically continues, the position IDs (and therefore the RoPE rotation) should also continue so that the relative positions remain consistent. However, if each chunk is treated as an independent sequence, then resetting to zero might be expected.

Since TorchTitan precomputes a RoPE cache, I’m trying to understand what the intended behavior is in this case.

[tianyu-l]

@francesco-bertolotti

If you want the two chunks still be parsed as a single sequence, obviously we should let position id continue. However, the attention part cannot model it -- the attention would consider them as independent, so only continuing the position ids doesn't make sense to me.

In reality, I think there are different modes of treating a sequence that cannot fit in the current row, e.g. padding, drop current example, etc.