DeepSeek-V4: eager attention and trainable FP8 grouped experts#7042
Conversation
deepseek_v4 ships a custom attention that is not compatible with the sdpa and flash paths, so add it to _EAGER_ONLY_PREFIXES to load with eager. Its fused experts load as FP8GroupedLinear, whose forward calls a grouped matmul kernel with no autograd formula, so loss.backward() fails during finetuning. Patch the forward to dequantize the frozen fp8 weight and run a differentiable grouped matmul while training, keeping the fused fp8 kernel for inference.
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7c3b676a01
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _FLASH_EXCLUDED_MODELS = ("gpt_oss",) | ||
| _EAGER_ONLY_PREFIXES = ("gemma3n",) | ||
| # deepseek_v4 ships sdpa/flash-incompatible custom attention; force eager. | ||
| _EAGER_ONLY_PREFIXES = ("gemma3n", "deepseek_v4") |
There was a problem hiding this comment.
For DeepSeek-V4 this only changes the resolver's default choice; if a caller passes requested_attn_implementation="sdpa" or "flash_attention_2", resolve_attention_implementation later overwrites the eager default with that requested backend because deepseek_v4 is not in the SDPA/flash exclusion checks. In contexts that forward an explicit attention kwarg to this helper, the model still reaches the sdpa/flash-incompatible attention path this change is trying to avoid, so the DeepSeek-V4 guard needs to participate in the final override/exclusion logic too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ed7ddd0: added deepseek_v4 to _SDPA_EXCLUDED_MODELS and _FLASH_EXCLUDED_MODELS, so an explicit attn_implementation of sdpa or flash downgrades to eager instead of reaching the unsupported path. Confirmed from_pretrained(attn_implementation=sdpa) now loads as eager.
| if self.weight.element_size() > 1 or not torch.is_grad_enabled(): | ||
| return _fp8_grouped_forward_orig(self, x) | ||
| hidden_dim = x.shape[-1] | ||
| W = weight_dequant(self.weight, self.weight_scale_inv.float()).to(x.dtype) |
There was a problem hiding this comment.
When an FP8GroupedLinear has been tensor-parallelized, Transformers stores weight/weight_scale_inv as distributed tensors and its original FP8 path unwraps them with to_local() before reshaping and calling the grouped matmul. This training replacement bypasses that unwrap, so TP DeepSeek-V4 finetuning feeds a distributed parameter into the local Triton dequantizer/view path and can fail or use the wrong shard metadata; unwrap both tensors to their local shards before weight_dequant.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Folded into ed7ddd0: the new autograd Function unwraps DTensor shards via to_local before dequant, matching the original fp8 path. unsloth does not tensor-parallelize these modules for finetuning today, so this is defensive, but it removes the divergence you noted.
| W = weight_dequant(self.weight, self.weight_scale_inv.float()).to(x.dtype) | ||
| w = W.view(self.n_groups, -1, hidden_dim).transpose(1, 2) | ||
| xg = x.reshape(-1, self.n_groups, hidden_dim).transpose(0, 1) | ||
| y = torch.bmm(xg, w).transpose(0, 1).reshape(*x.shape[:-2], self.n_groups, -1) |
There was a problem hiding this comment.
When training without gradient checkpointing, this torch.bmm makes autograd retain w, a bf16 dequantized copy of the full grouped projection weight, for every FP8GroupedLinear call until backward. DeepSeek-V4 has this grouped projection in each decoder block, so Pro-sized layers can add many GB of saved tensors and OOM; use a custom autograd path that saves the fp8 weight/scale and dequantizes during backward instead of letting bmm save the bf16 weight.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ed7ddd0: replaced the training bmm with a custom autograd Function that saves only the fp8 weight and scale rather than the bf16 dequantized copy, so no dequantized grouped weight is retained per layer; backward re-dequantizes. Bit-exact forward and grad with the previous path.
Add deepseek_v4 to _SDPA_EXCLUDED_MODELS and _FLASH_EXCLUDED_MODELS so an explicit attn_implementation=sdpa/flash request downgrades to eager instead of raising (the model has no sdpa/flash kernel), matching the eager-only default. Replace the FP8GroupedLinear training bmm with a custom autograd Function that saves only the fp8 weight + scale rather than a full bf16 dequantized copy, so no dequantized grouped weight is retained per layer, and unwrap tensor-parallel shards before dequant. Bit-exact forward and grad with the previous path.
|
@codex review |
1 similar comment
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ed7ddd0bd7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def forward(ctx, x, weight, scale_inv, n_groups, bias): | ||
| weight, scale_inv = _fp8_to_local(weight), _fp8_to_local(scale_inv) | ||
| hidden = x.shape[-1] | ||
| W = weight_dequant(weight, scale_inv.float()).to(x.dtype) |
There was a problem hiding this comment.
When an FP8GroupedLinear is quantized with a non-128 or rectangular weight_block_size, this training path calls weight_dequant without the layer's block size, so it falls back to the hard-coded 128x128 tiling. The existing _blockwise_weight_dequant_any_shape helper in this file handles exactly these cases because the default dequant can mis-index scales; here the custom backward path will train against weights dequantized differently from the original fused grouped matmul. Pass self.block_size into the custom function and use the block-size-aware helper instead of assuming 128.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 24c14b8: the function now dequantizes with the layer's own block_size via _blockwise_weight_dequant_any_shape, so non-128 or rectangular fp8 blocks are scaled correctly instead of assuming 128x128.
| return grad_x, None, None, None, grad_bias | ||
|
|
||
| def _fp8_grouped_forward(self, x): | ||
| if self.weight.element_size() > 1 or not torch.is_grad_enabled(): |
There was a problem hiding this comment.
When this layer is inside a reentrant gradient-checkpointed segment, the first training forward runs with grad disabled, so this condition takes the original fused grouped-matmul path, while the backward recomputation runs with grad enabled and takes the dequantized bmm path. Checkpointing assumes the recomputed forward is the same function as the one that produced the loss; using different FP8 kernels here means gradients are for different activations in checkpointed finetuning. Gate on training/checkpoint intent rather than torch.is_grad_enabled() alone, or make both paths use the same math during training.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 24c14b8: the training path now gates on self.training instead of torch.is_grad_enabled(), so a gradient-checkpointed segment runs the same bmm math in both its no-grad forward and its grad recompute (eval still uses the fused kernel). Verified the two passes are now bit-identical.
Gate the differentiable training path on self.training rather than torch.is_grad_enabled(), so a gradient-checkpointed segment runs the same bmm math in its no-grad forward and its grad recompute instead of mixing the fused fp8 kernel with bmm. Dequantize with the layer's own block_size via _blockwise_weight_dequant_any_shape so non-128 or rectangular fp8 blocks are scaled correctly instead of assuming 128x128.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Bravo. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Already looking forward to the next diff. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Summary
Finetuning DeepSeek-V4 (
deepseek_v4) withFastModelcurrently fails in two places. This fixes both so the model loads and trains.deepseek_v4ships a custom attention that is not compatible with the sdpa or flash paths, so it must run eager.loss.backward()errors. Transformers stores the fused experts asFP8GroupedLinear, whoseforwardcalls a grouped matmul kernel that has no autograd formula, so backward fails during finetuning.Changes
unsloth/models/_utils.py: adddeepseek_v4to_EAGER_ONLY_PREFIXESso it loads with eager attention.unsloth/kernels/fp8.py: patchFP8GroupedLinear.forwardso that during training it dequantizes the frozen fp8 weight and runs a differentiable grouped matmul, while inference keeps the original fused fp8 kernel. Gated ontorch.is_grad_enabled()and a fp8 weight, so the inference path is unchanged.Why bmm on the dequantized weight
torch._grouped_mm: forward matches exactly, but its backward raisesInvalid strides/sizeson the grouped-expert shapes, and backward is exactly what we need.torch._scaled_grouped_mm: needs a transposed operand and fp8-quantized activations, which adds quantization error and complexity.torch.bmmis exact at fp8 precision, has a finite backward, and works on all shapes. Only the few grouped-linear modules materialize in bf16, so the overhead is negligible.Testing
On
trl-internal-testing/tiny-DeepseekV4ForCausalLM:FP8GroupedLinearparity: dequant + bmm matches the fused fp8 kernel to 2.7e-2 relative (fp8 rounding), with finite gradients.Scope is loading and finetuning. Saving an fp8 checkpoint merged to 16bit is a separate path and is not changed here.