When doing inference with prompts of different length, having an attention mask is super useful so that model outputs are invariant to padding/batch size. Normally I expect that the model's answer to a question shouldn't change if I add pad tokens or batch it with another question, which is true for normal (autoregressive) LLMs.
For LLaDA, while the authors have recently added support for an attention_mask argument in model.forward(), which correctly masks out any padding tokens from the attention computation, there is still an issue with RoPE that causes the output to depend on batching/padding.
Since the RoPE module uses inferred positions based on only the tensor shape, adding pad tokens on the left (even with a proper attention mask) will shift all the subsequent positions, causing the qk vectors to change.
I think the RoPE module should accept a position_ids argument that allows it to ignore any pad tokens. This along with the attention mask will make the behavior deterministic regardless of padding.
When doing inference with prompts of different length, having an attention mask is super useful so that model outputs are invariant to padding/batch size. Normally I expect that the model's answer to a question shouldn't change if I add pad tokens or batch it with another question, which is true for normal (autoregressive) LLMs.
For LLaDA, while the authors have recently added support for an attention_mask argument in model.forward(), which correctly masks out any padding tokens from the attention computation, there is still an issue with RoPE that causes the output to depend on batching/padding.
Since the RoPE module uses inferred positions based on only the tensor shape, adding pad tokens on the left (even with a proper attention mask) will shift all the subsequent positions, causing the qk vectors to change.
I think the RoPE module should accept a position_ids argument that allows it to ignore any pad tokens. This along with the attention mask will make the behavior deterministic regardless of padding.