From 1c27f05370c42d612e78ed80c62db841a31f489c Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 18 Aug 2026 14:10:06 -0700 Subject: [PATCH] fix: guard backward_dw() bias-grad backfill against a legitimately-empty bgrad TransformerEngineBaseModule.backward_dw() unconditionally assumes the bgrad popped from wgrad_store is a real gradient tensor whenever use_bias is True, and crashes with AttributeError: 'NoneType' object has no attribute 'to' if it isn't. In FP8 mode, grad_output_preprocess() always computes the bias gradient eagerly (via bgrad_quantize or an unfused sum), regardless of delay_wgrad_compute. That value flows through the module's ordinary backward() return into normal autograd, which sets bias.grad. If the training framework's own gradient-accumulation hook then consumes bias.grad and resets it to None before backward_dw() runs (a common pattern for frameworks that manage their own master-gradient buffers, e.g. to fold .grad into a separate main_grad and free it for reuse), backward_dw()'s `if bias_tensor.grad is None:` check is satisfied even though there is nothing new to backfill -- the real gradient was already correctly handled by the eager path. Meanwhile, in linear.py's wgrad-GEMM closure, `"bias": (bias if (grad_bias is None and not bwd_args.fp8) else None)` deliberately skips computing bias grad in the wgrad step whenever grad_bias was already set, which in FP8 mode it always is -- so bgrad popped from wgrad_store is legitimately empty in this case, and backward_dw() crashes trying to use it anyway. Guard on bgrad actually holding data before assigning, matching the existing (grad_bias is not None and grad_bias.numel() != 0) pattern already used for the equivalent case in GroupedLinear.backward_dw(). Skipping is correct here rather than computing bgrad in the wgrad step too, since the latter would double-count a gradient that was already folded into the framework's own gradient buffer by the eager path. --- transformer_engine/pytorch/module/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index e9a65c3648..59a4d7e08a 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1961,7 +1961,7 @@ def backward_dw(self): if not self.fuse_wgrad_accumulation: weight_tensor = noop_cat(self._get_weight_tensors()) weight_tensor.grad = wgrad.to(weight_tensor.dtype) - if self.use_bias: + if self.use_bias and bgrad is not None and bgrad.numel() != 0: bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names]) if bias_tensor.grad is None: bias_tensor.grad = bgrad.to(bias_tensor.dtype)