huggingface / huggingface/candle

Qwen2: prepare_attention_mask is bidirectional (no causal triangle) when attn_mask provided

Open
#3,567 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

## Severity: CRITICAL — silently wrong outputs under batched eval / training

`Qwen2Model::prepare_attention_mask` at `candle-transformers/src/models/qwen2.rs:326-339` constructs the attention mask using **only the padding signal**, with no causal triangle. When an `attn_mask` is provided to `Model::forward()`, the model becomes a bidirectional encoder — every non-padding position can attend to every other non-padding position, including **future tokens**.

**Why this is silent.** Inference of a single un-padded prompt hits the `None` branch in `forward()` and is correct. The bug only triggers under batched eval / training / any caller that passes `attn_mask` for variable-length padding. Symptom: eval `token_acc = 100%` because the model trivially "predicts" tokens it can already see.

**Verified present at candle main HEAD `1d7e9274591ef7eb101d965b71243a41994214d2`** (committed 2026-05-26).

## Reproducer

https://github.com/nxrobins/candle-bug-6-qwen2-causal-mask

Standalone crate pinned to commit `7c7a8c570e6a16ea24cc30a7501c8ffbbdb51680` (the "Bump candle version to 0.10.2" commit, since 0.10.2 isn't tagged in the repo). One-line repro:

\`\`\`bash
git clone https://github.com/nxrobins/candle-bug-6-qwen2-causal-mask && cd candle-bug-6-qwen2-causal-mask && cargo run --release
# exit 1: BUG PRESENT (6 zeros above the diagonal of a 4×4 mask)
# exit 0: bug appears fixed
\`\`\`

Since `prepare_attention_mask` is `pub(crate)`, the reproducer replicates the upstream function byte-for-byte with attribution and ships a positive-control corrected implementation. Both \`cargo run\` and \`cargo test\` (2 tests) pass; no CUDA/Metal required.

## Suggested fix

Combine the padding mask with a lower-triangular causal mask via elementwise AND before converting to the \`0 / -inf\` attention bias:

\`\`\`rust
fn prepare_attention_mask(&self, attn_mask: &Tensor) -> Result {
let (b_sz, sql_len) = attn_mask.dims2()?;
let mut mask: Vec = vec![];
for b in 0..b_sz {
mask.push(attn_mask.i((b, ..))?.expand((1, 1, sql_len, sql_len))?);
}
let padding_mask = Tensor::cat(&mask, 0)?;

// Lower-triangular causal mask.
let causal: Vec = (0..sql_len)
.flat_map(|i| (0..sql_len).map(move |j| if j <= i { 1.0 } else { 0.0 }))
.collect();
let causal_2d = Tensor::from_slice(&causal, (sql_len, sql_len), &self.device)?;
let causal_bool = causal_2d
.reshape((1, 1, sql_len, sql_len))?
.broadcast_as(padding_mask.shape())?
.to_dtype(padding_mask.dtype())?;

// A cell is allowed iff BOTH padding and causal allow it.
let combined = (padding_mask * causal_bool)?;

let on_true = combined.zeros_like()?.to_dtype(self.dtype)?;
let on_false = Tensor::new(f32::NEG_INFINITY, &self.device)?
.broadcast_as(combined.shape())?
.to_dtype(self.dtype)?;
combined.where_cond(&on_true, &on_false)
}
\`\`\`

This was applied locally to a vendored copy of \`qwen2.rs\` and verified end-to-end against a Python reference: \`token_acc\` returned to expected ~35% on Qwen2.5-0.5B (was 100% with the bug).

## Likely affected sibling models

Any model in \`candle-transformers/src/models/\` with a similar \`prepare_*_mask\` helper that doesn't include causality. \`git grep -lE \"fn prepare_(attention|padding)_mask\" candle-transformers/src/models/\` returns the audit list. Happy to audit + file per-model PRs if useful.

## Context

Surfaced during Phase 0.5b validation of an external project against a Python reference implementation. Full writeup with the four bugs that surfaced together: https://github.com/nxrobins/wave-agent/blob/main/docs/upstream-candle-bugs.md

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with candle-transformers/src/models/qwen2.rs:326-339 and compare prepare_attention_mask with the standalone reproducer and its Python reference. Run cargo test or cargo run --release in candle-bug-6-qwen-causal-mask; the work is done when the reproducer exits 0 and batched attention no longer permits future-token access.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
70/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.