fix: guard backward_dw() bias-grad backfill against a legitimately-empty bgrad

[nvegesna-netizen]

Summary

TransformerEngineBaseModule.backward_dw() crashes with
AttributeError: 'NoneType' object has no attribute 'to' when delay_wgrad_compute
is enabled together with FP8 and a biased Linear/LayerNormLinear layer, under a
training setup where something external resets bias.grad back to None between the
module's ordinary backward() and the later backward_dw() call (e.g. a
gradient-accumulation hook that folds .grad into its own master-gradient buffer and
clears it for reuse — a common pattern in large-scale training frameworks).

Root cause

1. 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.
2. If something external then consumes bias.grad and resets it to None (as
   described above) before backward_dw() runs, 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.
3. Separately, linear.py's wgrad-GEMM closure builds its kwargs as
   "bias": (bias if (grad_bias is None and not bwd_args.fp8) else None) — in FP8
   mode this is always None, since step 1 already computed the real gradient. So the
   wgrad GEMM never computes a bias grad in FP8 mode, and bgrad popped from
   wgrad_store in backward_dw() is legitimately empty.
4. backward_dw() doesn't distinguish "nothing to backfill because it was already
   handled" from "genuinely missing" — it unconditionally tries bgrad.to(...),
   crashing on the empty value from step 3.

This only affects FP8 (non-FP8 layers genuinely defer bias-grad computation to the
wgrad GEMM — grad_output_preprocess's non-FP8 branch returns grad_bias=None on
purpose there, so bgrad is real in that path and backward_dw() correctly applies
it; this fix doesn't touch that path, since bgrad is non-empty there).

Fix

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 situation 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 caller's own gradient buffer by the eager
path — a silent numerics bug, strictly worse than the crash it would replace.

Testing

I don't have GPU access in the environment I used to prepare this change, so I could
not run TE's own test suite directly against it. What I have validated, on real
GPU hardware, training a MoE-style model with FP8 delayed scaling and a training
framework that performs the external-.grad-reset pattern described above (an
interleaved-pipeline-parallel schedule invoking backward_dw() after each
microbatch's backward pass, with a biased QKV projection):

• Confirmed the crash reproduces exactly as described, at the bgrad.to(...) line,
  on the unmodified module.
• Applied this exact one-line fix and confirmed training proceeds normally past that
  point with the same configuration.
• Separately ablation-tested with only an unrelated, independent fix applied (to a
  different module's analogous bug) and confirmed this specific crash still
  reproduces without this change — i.e. this fix is independently necessary, not
  incidental to something else.

On coverage in this repo's own test suite: tests/pytorch/test_numerics.py has
test_linear_accuracy_delay_wgrad_compute, which already exercises delay_wgrad_compute
with bias=True, but it doesn't currently reproduce this bug because (a) it's
parametrized only over torch.float32/float16/bfloat16, never FP8, and (b) its
harness calls loss.backward() immediately followed by block.backward_dw() with
nothing in between resetting bias.grad — so bias.grad is simply left set from the
ordinary backward pass, and backward_dw()'s if bias_tensor.grad is None: guard
skips the crashing branch entirely in that harness. Reproducing this bug in that test
would need an FP8-enabled parametrization plus an explicit bias.grad = None between
backward() and backward_dw() to model the external-consumer pattern. I'm not
proposing that test addition myself in this PR since I can't execute it locally to
confirm it's correct, but wanted to flag the precise gap for maintainers, since it's a
straightforward extension of the existing harness if useful.

[greptile-apps]

Greptile Summary

The PR prevents delayed weight-gradient backfill from dereferencing an absent or empty bias gradient after FP8 bias-gradient handling has already occurred.

• Guards bias-gradient assignment on bgrad being non-null and non-empty.
• Preserves the existing non-FP8 delayed bias-gradient backfill path.

Confidence Score: 5/5

The PR appears safe to merge, with the new guard correctly handling the legitimately absent FP8 bias-gradient value.

Current delayed-wgrad producers supply either a real bias-gradient tensor or None, and the changed condition preserves real-gradient assignment while safely skipping the absent value that previously caused the crash.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Adds a narrowly scoped validity guard before delayed bias-gradient assignment; no correctness issue was identified in the reachable producer paths.

_{Reviews (1): Last reviewed commit: "fix: guard backward_dw() bias-grad backf..." | Re-trigger Greptile}