From 4621b7a282f81a6f7846819c6cfaa01a104c44ab Mon Sep 17 00:00:00 2001 From: Rudra Khunti Date: Tue, 4 Aug 2026 16:38:13 +0530 Subject: [PATCH] Fix training crash in addmm_act by falling back to unfused ops when grad is enabled The fused _addmm_activation path detaches weights and casts to bf16, so it only works under no_grad. Since the SAM 3.1 update, vitdet's Mlp.forward() calls addmm_act unconditionally, which makes any training/fine-tuning run (e.g. the documented Roboflow configs) fail with 'ValueError: Expected grad to be disabled.' (#610). Fall back to the equivalent standard linear+activation ops when gradients are enabled; inference keeps the fused fast path unchanged. --- sam3/perflib/fused.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sam3/perflib/fused.py b/sam3/perflib/fused.py index 6800cca64..4f7455657 100644 --- a/sam3/perflib/fused.py +++ b/sam3/perflib/fused.py @@ -9,7 +9,14 @@ def addmm_act(activation, linear, mat1): if torch.is_grad_enabled(): - raise ValueError("Expected grad to be disabled.") + # Training path: the fused kernel below is inference-only (detached + # weights, bf16 cast), so fall back to standard autograd-friendly ops. + y = linear(mat1) + if activation in [torch.nn.functional.relu, torch.nn.ReLU]: + return torch.nn.functional.relu(y) + if activation in [torch.nn.functional.gelu, torch.nn.GELU]: + return torch.nn.functional.gelu(y) + raise ValueError(f"Unexpected activation {activation}") self = linear.bias.detach() mat2 = linear.weight.detach() self = self.to(torch.bfloat16)