fix(ragged_dot): store backward gradients in fp32 to prevent bf16 weight-grad overflow - #1
Open
emergenz wants to merge 1 commit into
Open
fix(ragged_dot): store backward gradients in fp32 to prevent bf16 weight-grad overflow#1emergenz wants to merge 1 commit into
emergenz wants to merge 1 commit into
Conversation
The custom VJP in `ragged_dot/base.py` computed the two backward dots with `preferred_element_type=lhs.dtype` (dlhs) and `preferred_element_type=rhs.dtype` (drhs), i.e. it stored the gradients at the operand dtype. With bf16 operands the fp32 accumulator is downcast to bf16 on store, so a weight gradient that lands in (bf16_max=3.3895e38, fp32_max=3.4028e38] overflows to `inf` even when the bf16 forward is finite. Store the backward gradients in a wider dtype (fp32 by default, via the new defaulted `grad_dtype` param + `_backward_grad_dtype` helper) while keeping the bf16 operands, so the tensor-core multiply is unchanged but the gradient is no longer rounded through bf16. This matches the accumulation dtype of the forward and of jax.lax.ragged_dot's own VJP. Change is additive/defaulted; all base.vjp callers (PallasTriton, Mosaic GPU/TPU, RaggedDot(vjp=base.vjp)) are covered. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.
Problem
RaggedDot's custom VJP (base.vjp) stored the backward gradients at the operand dtype — the two backward dots usedpreferred_element_type=lhs.dtype/rhs.dtype. With bf16 operands, the down_proj weight-gradient can land in the band (bf16_max 3.3895e38, fp32_max 3.4028e38] and overflow toinfon the bf16 store while the forward stays finite — a "forward fine, backward NaN" failure in bf16 (MoE) training. The accumulator was already fp32; only the gradient storage dtype was the problem.Fix
Backward-only, single file (
_src/ops/ragged_dot/base.py). Adds a defaultedgrad_dtypeparam + a_backward_grad_dtypehelper; the two backward dots now usepreferred_element_type = result_type(operand_dtype, float32)(bf16→fp32, fp32→fp32; explicit override honored for future fp8 work). The forward is untouched — the bf16 tensor-core multiply is preserved; only the gradient storage/output dtype widens to fp32.Additive/defaulted; covers all
base.vjpcallers (PallasTriton, Mosaic GPU/TPU viafunctools.partial, explicit-vjp xla) and bothragged_dot/ragged_dot_general. The bareRaggedDot()/implementation="xla"path uses JAX's native autodiff (notbase.vjp) and is intentionally out of scope (matchesjax.lax.ragged_dot). Returning fp32 cotangents for bf16 operands is permitted by JAX'scustom_vjp(_temporary_dtype_exception).Verification (A100-40GB, JAX 0.9.2)
inf).base_test.py557/557;pallas_triton_test.py276 incl. all 24test_vjp(3 pre-existing forward-only failures fail identically on unmodified tokamax). Forward output bit-identical before/after.Impact
Lets downstream (omegalax) drop its fp32 operand-upcast workaround — recovering ~1.6–2× on bf16 MoE expert GEMMs by keeping the bf16 multiply while getting safe fp32 gradient storage. Downstream must accept fp32 gradients for bf16
ragged_dotoperands (a no-op widening for fp32-master-weight setups).fp8-quantization-aware
ragged_dotis tracked separately — there is no fp8×fp8 grouped kernel in tokamax on any architecture today, so that's a distinct effort (weight-only fwd + XLA-simulated backward now, or a DeepSeek/DeepGEMM-style fp8×fp8 grouped kernel as a larger project).