Skip to content

bug: Gemma 4 26B A4B standard-4-bit deterministically emits a wrong token (Middle War for Middle Ages) in long multi-turn context #933

Closed

Description

@inureyes

Summary

After #910 (the qmm_sm80 async race fix) made 4-bit CUDA decode deterministic, gemma-4-26b-a4b-it-4bit still deterministically emits a semantically wrong token at one position in a specific long multi-turn context: the wine paragraph of the 5-paragraph repro reads "During the Middle War, monasteries became vital centers for preserving winemaking" where it should read "Middle Ages". The rest of the paragraph is coherent. This is NOT the #910 race (see evidence) and NOT random garbage (no foreign-script, fully deterministic). It is a standard-4-bit (RTN/affine) quantization-quality artifact: the quantization-aware-trained variant does not exhibit it.

Reproduction

printf '%s\n' \
"Remember this secret code: PARETO-7391. Now write a detailed 160-word paragraph about the history of coffee." \
"Now write a detailed 160-word paragraph about the history of tea." \
"Write a detailed 160-word paragraph about the history of chocolate." \
"Write a detailed 160-word paragraph about the history of wine." \
"Write a detailed 160-word paragraph about the history of bread." \
"/bye" | ./target/release/mlxcel run <path>/gemma-4-26b-a4b-it-4bit -t 0 -n 256

Wine paragraph contains "During the Middle War, monasteries...". Reproduces 4/4 runs, byte-identical.

Evidence that this is quantization quality, not a runtime bug

Root-cause hypothesis (to confirm)

Standard 4-bit (affine, group_size 64) quantization degrades one near-decision in a long-context regime enough to flip the greedy argmax. The #910 investigation measured (pre-fix, on cleaner contexts) that "Ages" normally wins by 14-17 logits, so the flip is context-specific; the likely sensitive component is the 262k-token embedding/unembedding (logit head) and/or specific decoder layers under 4-bit. The QAT variant's immunity supports a weight-precision cause.

Investigation and fix directions

  1. Quantify the margin post-bug: Gemma 4 26B A4B temp-0 decode is non-deterministic run-to-run and corrupts near-tied tokens independently of the fused MoE kernel #910: dump top-k logits at the exact "During the Middle ___" decode step in the multi-turn context, for standard-4bit vs qat-4bit, to measure how narrowly "War" beats "Ages" and whether the logit head is the dominant contributor.
  2. Test targeted higher-precision: does keeping the unembedding / logit head (or the embedding) at higher precision (e.g. 8-bit or bf16) eliminate the flip? If a small, bounded precision bump on one sensitive tensor fixes it at acceptable memory/perf cost, that is the fix (and a general 4-bit quality win, not gemma4-specific).
  3. Confirm QAT robustness across a few more prompts (is it a reliable mitigation to recommend).
  4. If no bounded code fix exists (i.e. it is inherent RTN 4-bit degradation), conclude with evidence and document mitigations (use the QAT checkpoint, or an 8-bit / higher-precision-head quantization) rather than shipping a speculative change.

Acceptance criteria

  • Either: a bounded, verified code/quantization change (e.g. higher-precision logit head by default, or a documented mixed-precision option) that makes the standard-4bit repro emit "Middle Ages" (matching QAT / higher precision) without unacceptable memory/perf regression and with no regression on other models; the 5-paragraph repro is clean and deterministic.
  • Or: an evidence-backed determination that this is inherent standard-4-bit RTN degradation, with a documented mitigation (QAT / 8-bit / higher-precision head) and the logit-margin data captured on the issue.

Priority / impact

Low. Deterministic single wrong word (not garbage, not non-deterministic) on the standard-4-bit checkpoint only; a working mitigation (the QAT variant) already exists. Much milder than the pre-#910 corruption.

Environment / related

Activity

  1. inureyes commented on Jul 27, 2026

    @inureyes
    MemberAuthor

    Investigation result: no bounded compute-precision fix; inherent standard-4-bit degradation (can-not-fix)

    Confirmed on current main (4fd625b, post-#910) on GB10. The 5-paragraph repro deterministically emits "During the Middle War" on gemma-4-26b-a4b-it-4bit, while the bread paragraph in the same run correctly emits "During the Middle Ages". This is standard-4-bit (RTN/affine) quantization quality, not a runtime bug, and the prime hypothesis (that the logit-head accumulation precision is the lever, analogous to #886) is disproven by direct measurement.

    Where the argmax is decided and what the margin is

    The failing token is the projection hidden(2816) -> vocab(262144) at the wine "During the Middle ___" decode step. Instrumenting the tied-embedding logit head to compute BOTH the current bf16 head and an f32-accumulated head from the IDENTICAL bf16 hidden state at that exact step (token ids: " War"=5289, " Ages"=52901; values are recovered pre-softcap logits, cap=30):

    • bf16 head (current): argmax = " War". Raw logits: War = -18.50, Ages = -22.50. Margin War over Ages = 4.00.
    • f32-accumulated head (candidate fix): argmax = still " War". Raw logits: War = -18.455, Ages = -22.498. Margin = 4.04.

    The bf16-vs-f32 accumulation difference on these logits is about 0.045 (War) and 0.002 (Ages), roughly 100x smaller than the ~4.0-logit margin, and it does not change the argmax. For contrast, at the bread "Middle Ages" step the model is confident (Ages raw ~+21.6, next candidate " Age" ~+1.35, " War" ~-5.5); the wine long-context regime is where every candidate collapses into a low clustered band (top-1 raw only about -18.5) and " War" edges out " Ages".

    Enabling the f32-accumulated logit head for the entire run (which perturbs the whole greedy trajectory via the butterfly effect) still emits "During the Middle War".

    The accumulation site (for the record)

    The CUDA qmv kernel (src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu) does accumulate the per-thread K-partials in the activation dtype for bits < 8 (conditional_t<(bits >= 8), float, T>), i.e. bf16 for this 4-bit head, while bits >= 8 accumulates in f32. So there is a genuine bf16-accumulation site, and forcing it to f32 is a real change (note: MLX's CUDA quantized_matmul output dtype tracks the scales/biases dtype, so an honest f32 test must upcast scales+biases, not just the activation). But the measured effect of that change at the failing decision is negligible relative to the margin, so it is not a fix. The ~4-logit gap is baked into the bf16 hidden state and the 4-bit weights before the projection; f32-accumulating the final projection moves it by <0.05.

    Conclusion and mitigations

    This is inherent standard-4-bit (RTN/affine, group_size 64) weight degradation surfacing only in the long accumulated multi-turn context, not a bounded compute/accumulation-precision defect. No speculative code change is being shipped (same posture as #817 / #824).

    Confirmed mitigations:

    • Use gemma-4-26b-a4b-it-qat-4bit (quantization-aware-trained): does not exhibit the flip on the identical repro (only correct "Middle Ages" / "Middle East").
    • Use a higher-bit checkpoint (e.g. 8-bit): 8-bit weights carry less error AND the qmv already accumulates 8-bit in f32.

    Priority stays low: deterministic single wrong word (not garbage, not non-deterministic) on the standard-4-bit checkpoint only, with a working mitigation already available. Reopen if a product decision is made to change defaults (e.g. a higher-precision logit head by default, or f32 accumulation across the whole 4-bit decode path), but neither is a bounded logit-head change and the QAT mitigation already covers the case.

  2. added and removed on Jul 27, 2026
  3. inureyes commented on Jul 27, 2026

    @inureyes
    MemberAuthor

    Additional data point from a release-regression sweep (post-#910, binary 4fd625b): the stray-token artifact is not limited to the wine/"Middle War" context. On a simple prompt ("List three interesting facts about Mars..."), standard-4bit gemma-4-26b-a4b deterministically emits "...from New York City to Los जाते Los Angeles" (a stray Devanagari token) while the rest of the output is fully correct. It is deterministic (byte-identical across runs, matches the CUDA_LAUNCH_BLOCKING=1 oracle), and the QAT variant (gemma-4-26b-a4b-it-qat-4bit) is clean on the same prompt. This confirms the wontfix diagnosis: inherent standard-4bit RTN degradation of this checkpoint (occasional deterministic stray tokens across contexts), mitigation is the QAT checkpoint. In the same sweep, 13 other models (dense/MoE/VLM/bf16/mxfp4, 4-bit incl. mixtral, qwen3-30b-a3b, gpt-oss-20b) were clean and deterministic, so this is isolated to the standard-4bit gemma-4-26b-a4b weights.

  4. self-assigned this
    on Aug 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

priority:lowLow prioritystatus:wontfixThis will not be worked ontype:bugBug fixes, error corrections, or issue resolutions

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions