From f899834e58a735a7fa7d5341bf944476b32a835e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 19:33:05 -0700 Subject: [PATCH] DeepSeek-V4: eager attention and trainable FP8 grouped experts (#7042) * DeepSeek-V4: eager attention and trainable FP8 grouped experts 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. * DeepSeek-V4: exclude sdpa/flash and stream fp8 grouped backward 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. * FP8 grouped: consistent checkpointing math and block-size-aware dequant 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/kernels/fp8.py | 68 ++++++++++++++++++++++++++++++++++++++++ unsloth/models/_utils.py | 8 +++-- 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 4efc4bd5d3..7a57f91ce1 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -42,6 +42,11 @@ except: "Unsloth: FP8 models need importing FP8Linear from `transformers.integrations.finegrained_fp8` but we don't see it." ) +try: + from transformers.integrations.finegrained_fp8 import FP8GroupedLinear +except: + FP8GroupedLinear = None + try: from transformers.integrations.fbgemm_fp8 import FbgemmFp8Linear except: @@ -688,3 +693,66 @@ if FbgemmFp8Linear is not None: FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, "weight_scale") if FP8Linear is not None: FP8Linear.forward = module_forward_patch(fp8_block_quant_linear, "weight_scale_inv") + +# FP8GroupedLinear's fused grouped matmul has no autograd formula, so training +# backward fails. In training, use a custom autograd Function: dequant the frozen +# fp8 weight for a differentiable bmm, saving only the fp8 weight + scale and +# unwrapping TP shards; eval keeps the fused kernel. Gate on self.training (not +# is_grad_enabled) so the grad-checkpoint no-grad forward and its recompute match. +if FP8GroupedLinear is not None: + _fp8_grouped_forward_orig = FP8GroupedLinear.forward + + def _fp8_to_local(t): + dt = getattr(getattr(torch, "distributed", None), "tensor", None) + DTensor = getattr(dt, "DTensor", None) if dt is not None else None + return t.to_local() if DTensor is not None and isinstance(t, DTensor) else t + + def _fp8_grouped_dequant(weight, scale_inv, block_size, dtype): + # Honor the layer's block size; weight_dequant would assume 128 and mis-scale. + if block_size is not None and len(block_size) == 2: + return _blockwise_weight_dequant_any_shape(weight, scale_inv.float(), block_size, dtype) + return weight_dequant(weight, scale_inv.float()).to(dtype) + + class _FP8GroupedMM(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, scale_inv, n_groups, block_size, bias): + weight, scale_inv = _fp8_to_local(weight), _fp8_to_local(scale_inv) + hidden = x.shape[-1] + W = _fp8_grouped_dequant(weight, scale_inv, block_size, x.dtype) + out_per = W.shape[0] // n_groups + xg = x.reshape(-1, n_groups, hidden).transpose(0, 1) + y = torch.bmm(xg, W.view(n_groups, out_per, hidden).transpose(1, 2)) + y = y.transpose(0, 1).reshape(*x.shape[:-2], n_groups, out_per) + if bias is not None: + y = y + bias.view(n_groups, out_per) + ctx.save_for_backward(weight, scale_inv) + ctx.n_groups, ctx.out_per, ctx.x_shape = n_groups, out_per, x.shape + ctx.dtype, ctx.has_bias, ctx.block_size = x.dtype, bias is not None, block_size + return y + + @staticmethod + def backward(ctx, grad_y): + weight, scale_inv = ctx.saved_tensors + ng, out_per, hidden = ctx.n_groups, ctx.out_per, ctx.x_shape[-1] + W = _fp8_grouped_dequant(weight, scale_inv, ctx.block_size, ctx.dtype).view( + ng, out_per, hidden + ) + gy = grad_y.reshape(-1, ng, out_per).transpose(0, 1) + grad_x = torch.bmm(gy, W).transpose(0, 1).reshape(ctx.x_shape) + grad_bias = gy.sum(1).reshape(-1) if ctx.has_bias else None + return grad_x, None, None, None, None, grad_bias + + def _fp8_grouped_forward(self, x): + if self.weight.element_size() > 1 or not self.training: + return _fp8_grouped_forward_orig(self, x) + bias = self.bias if self.has_bias else None + return _FP8GroupedMM.apply( + x, + self.weight, + self.weight_scale_inv, + self.n_groups, + getattr(self, "block_size", None), + bias, + ) + + FP8GroupedLinear.forward = _fp8_grouped_forward diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index fa0e0b1c49..ae3b94ecc8 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -424,7 +424,7 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len # access on some GPU architectures (B200). Falls back to eager safely. _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") -_SDPA_EXCLUDED_MODELS = ("gpt_oss",) +_SDPA_EXCLUDED_MODELS = ("gpt_oss", "deepseek_v4") # The loader (loader.py) forces supports_sdpa=False for these because their bundled # SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor # them without a loader -> _utils import cycle (loader.py already imports from _utils @@ -437,8 +437,10 @@ DISABLE_SDPA_MODEL_NAMES = [ "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore "gpt_oss", ] -_FLASH_EXCLUDED_MODELS = ("gpt_oss",) -_EAGER_ONLY_PREFIXES = ("gemma3n",) +_FLASH_EXCLUDED_MODELS = ("gpt_oss", "deepseek_v4") +# deepseek_v4's custom attention is sdpa/flash-incompatible; force eager, and +# excluded above so an explicit sdpa/flash request cannot re-enable the crash. +_EAGER_ONLY_PREFIXES = ("gemma3n", "deepseek_v4") _FLASH_ATTENTION_MAX_HEAD_DIM = 256 _FLASH_ATTENTION_DISABLED_WARNED = set()