fix: clamp tile_size to dim when axis is smaller than tile_size - #1
Open
Insideyyy wants to merge 9 commits into
Open
fix: clamp tile_size to dim when axis is smaller than tile_size#1Insideyyy wants to merge 9 commits into
Insideyyy wants to merge 9 commits into
Conversation
When subchannel tile_size exceeds the axis dimension (e.g., tile_size=256 on an axis of size 16), fall back to per-tensor scale for that axis by clamping tile_size to dim. This avoids ValueError for operations like attention output projection where contraction axes can be small. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Insideyyy
force-pushed
the
fix-tile-size-fallback
branch
from
April 25, 2026 02:21
f83871a to
c6c7541
Compare
When multiple contraction axes exist (e.g. wgrad contracting on B,S), tile the last axis (S) for subchannel and make remaining contraction axes (B) channelwise. This ensures wgrad uses proper subchannel quantization instead of falling back to slow path. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
MLA's wkv_a layer has output dim 576 (kv_lora_rank 512 + qk_rope_head_dim 64), which is not divisible by tile_size 256. Instead of crashing with ValueError, gracefully degrade to per-tensor quantization for that axis with a warning. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Previously, convert_to() ignored noise_fn for floating-point qtypes (including float8_e5m2/e4m3fn), causing stochastic rounding to have no effect. This resulted in identical loss curves regardless of bwd_stochastic_rounding config. Now applies noise before the final astype() for FP8, matching the behavior of integer quantization. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…for FP8 1. Tracer leak fix: rng_key from make_rng was captured in a Python closure and stored in config as nondiff_argnum, causing the tracer to escape remat boundary. Pass rng_key as explicit traced JAX argument to custom_vjp instead. 2. FP8 stochastic rounding: scale noise by ulp since FP8 values are non-uniform. Without this, noise in (-0.5, 0.5) had no effect. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Adding ulp-scaled noise for stochastic rounding can push values beyond the representable FP8 range, producing inf (e.g. e5m2 qmax=57344, max ulp=8192, so values can exceed range → inf). Re-clip after noise. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ification Channelwise noise shared along the contraction axis causes coherent error accumulation in the dot product, amplifying gradient variance by O(K). Use elementwise noise (channelwise_noise_axes=None) for backward SR instead. Also add comment clarifying the ULP clamp logic in numerics.py. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The previous "add noise then RNE" approach was fundamentally broken for non-uniform FP8 types. Noise scaled by ULP at the current exponent could push values across exponent boundaries where step sizes differ, causing rounding to wrong grid points (e.g., x=1.0 in e5m2 would incorrectly round down to 0.875 with P=0.25 instead of staying at 1.0). Replace with mathematically correct floor/ceil stochastic rounding: 1. Find floor_fp8 (largest FP8 value <= x) and ceil_fp8 (smallest >= x) 2. Compute P(ceil) = (x - floor) / (ceil - floor) 3. Stochastically choose between floor and ceil based on this probability This produces unbiased rounding for any non-uniform quantization scheme. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
HowToQuantize already carries a noise_fn field and quantize_with_scale_zero_point already forwards it to numerics.convert_to, but quantize_api (the public qpl.quantize entry) didn't accept it — so callers outside dot_general_qt had no way to get SR into their bwd quantization. This unblocks MaxText's MoE GMM bwd path: megablox._gmm_bwd can now build a noise_fn from QtRule.bwd_stochastic_rounding and pass it directly to qpl.quantize, mirroring how dot_general_qt handles SR internally. No behavior change when noise_fn is omitted.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When subchannel tile_size exceeds the axis dimension (e.g., tile_size=256 on an axis of size 16), fall back to per-tensor scale for that axis by clamping tile_size to dim. This avoids ValueError for operations like attention output projection where contraction axes can be small.
Fixes #<issue_number_goes_here>