Fix 4-bit split QKV projection#1502
Conversation
|
@jlarson4 ready for review |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for this fix! I verified it locally. The unit-tier, fake-bnb approach for CPU-safe coverage is a nice touch.
There are a couple organization changes outlined below that we should address before merging, but altogether the solution is sound.
Outside of the current edits in this PR, it would be appreciated if you could update tests/QUARANTINES.md.
The new test moves the two quarantined entries in tests/unit/components/test_attention.py (cited there as :48 and :83) down by ~72 lines. Please refresh those line references.
|
|
||
| return q, k, v | ||
|
|
||
| def _project_4bit_qkv( |
There was a problem hiding this comment.
The 4D path use n_heads times the necessary compute and memory – please dequantize once and reuse complex_attn_linear instead.
_project_4bit_qkv projects every input head through all heads' weights and then keeps only the diagonal. On a 4-bit Llama-7B (32 heads) at pos=2048, that's a ~512 MiB fp16 (~1 GiB fp32) transient per projection, three times per layer, of which only the ~16 MiB diagonal survives. This is a real OOM risk on the memory-constrained setups people use 4bit models to avoid.
Since matmul_4bit dequantizes internally anyway, dequantizing the weight lets the split path share the already-tested einsum instead of maintaining a parallel implementation. Divergence between these two paths is how #737 happened in the first place.
There was a problem hiding this comment.
Additionally, if you could possibly add a docstring comment to this function?
| @@ -446,59 +446,20 @@ def calculate_qkv_matrices( | |||
| ) | |||
| if self.cfg.load_in_4bit: | |||
| W_Q_4bit = cast(Params4bit, self.W_Q) | |||
There was a problem hiding this comment.
W_Q currently gets a static cast while W_K/W_V get the runtime isinstance + ValueError. A non-quantized W_Q under load_in_4bit=True dies with an opaque AttributeError.
Let's move the isinstance guard from the W_K/W_V cases into the helper and drop the Q/K/V asymmetry. Checking isinstance(weight, Params4bit) inside _project_4bit_qkv gives all three weights the clear error and collapses the three near-identical call-site blocks.
|
@abhinav-bellapu Update looks good, once the CI passes I will merge |
|
@jlarson4 CI looks good! |
2 participants
Fixes #737.
Summary
[batch, pos, d_model]Q/K/V inputs.[batch, pos, head, d_model]inputs, dequantize the 4-bit weight once and reusecomplex_attn_linearfor the head-wise projection.Params4bitvalidation in the shared projection helper so Q, K, and V report the same clear error.Why
The old 4-bit path reshaped split Q/K/V projections incorrectly, producing the shape error in #737. Projecting every input head through every output head before selecting the diagonal would fix the shape but introduce an
n_heads-fold compute and memory cost. Dequantizing once and sharing the established head-wise projection path preserves the intended semantics without that transient.Validation
env UV_PYTHON=3.12 make formatenv UV_PYTHON=3.12 make check-formatuv run --python 3.12 pytest tests/unit/components/test_attention.py -q(15 passed, 4 skipped)uv run --python 3.12 mypy .(300 source files clean)env UV_PYTHON=3.12 make unit-test(3,027 passed, 33 skipped, 9 xfailed)