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.
This commit is contained in:
Daniel Han 2026-07-09 15:07:24 +00:00
parent b5aef63c03
commit 7c3b676a01
2 changed files with 29 additions and 1 deletions

View file

@ -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,25 @@ 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.forward calls a grouped matmul kernel with no autograd
# formula, so backward fails while training. For training, dequantize the frozen
# fp8 weight and run a differentiable grouped matmul; inference keeps the kernel.
# torch._grouped_mm matches but has no backward here; _scaled_grouped_mm needs
# fp8 activations, so bmm on the dequantized weight is exact and universal.
if FP8GroupedLinear is not None:
_fp8_grouped_forward_orig = FP8GroupedLinear.forward
def _fp8_grouped_forward(self, x):
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)
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)
if self.has_bias:
y = y + self.bias.view(self.n_groups, -1)
return y
FP8GroupedLinear.forward = _fp8_grouped_forward

View file

@ -438,7 +438,8 @@ DISABLE_SDPA_MODEL_NAMES = [
"gpt_oss",
]
_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")
_FLASH_ATTENTION_MAX_HEAD_DIM = 256
_FLASH_ATTENTION_DISABLED_WARNED = set()