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

[nvegesna-netizen]

Summary

TransformerEngineBaseModule.backward_dw() can crash with
AttributeError: 'NoneType' object has no attribute 'to' when delayed weight-gradient
computation is used with FP8 and a biased Linear or LayerNormLinear. The failure
occurs if an external gradient consumer moves the eagerly computed bias gradient into
its own accumulation buffer and clears bias.grad before backward_dw().

Root cause

In FP8 mode, grad_output_preprocess() computes the bias gradient eagerly. The
module's ordinary backward returns that gradient to autograd.

The delayed wgrad closure does not request another bias gradient in FP8 mode. Its
GEMM therefore returns bgrad=None, which is valid because the eager path already
produced the real bias gradient.

If an external hook consumes the eager gradient and clears bias.grad, the old
backward_dw() implementation enters its bias backfill branch and calls
bgrad.to(...) even though the deferred GEMM did not produce a bias gradient.

Current framework integrations that honor TE's skip_backward_post_hook contract may
defer their gradient hook until backward_dw() and avoid this ordering. The failing
state is still valid for an unconditional external gradient consumer, and it has been
observed in an integration where that deferral was not preserved.

Fix

Only backfill bias.grad when the delayed GEMM returned a nonempty bias gradient.
This matches the existing guard in GroupedLinear.backward_dw().

The non-FP8 path is unchanged. In that path the deferred GEMM returns a real bias
gradient, so the existing backfill still runs. In the FP8 failure path, skipping the
empty value preserves the eager gradient already consumed by the caller and avoids
double counting.

Testing

Added a focused GPU regression in tests/pytorch/test_numerics.py. It uses delayed
scaling FP8 and delayed wgrad computation, then installs a post-accumulate parameter
hook that moves the eager bias gradient into main_grad and clears bias.grad. The
test verifies that:

• the exact expected bias gradient is accumulated once
• backward_dw() completes without recreating or double counting the bias gradient
• delayed weight gradients and input gradients remain finite

Validation on one H100 with the pinned NeMo rc8 Transformer Engine build produced the
expected red and green result. The unpatched build failed at bgrad.to(...) with the
reported NoneType exception. The same test passed after applying this one-line fix.

All 24 existing H100 cases selected by
test_linear_accuracy_delay_wgrad_compute also passed after the patch. Those cases
cover the unaffected non-FP8 delayed-wgrad path across supported dtypes, batch sizes,
bias settings, and fused wgrad accumulation settings.

[greptile-apps]

Greptile Summary

The PR prevents delayed weight-gradient computation from dereferencing a legitimately absent FP8 bias gradient and adds regression coverage for frameworks that consume and clear bias.grad.

• Guards deferred bias-gradient backfill against None and empty tensors.
• Verifies the eager FP8 bias gradient remains preserved in an external master-gradient buffer.
• Confirms delayed weight-gradient computation still produces finite input and weight gradients.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	The new guard safely skips absent FP8 deferred bias gradients while preserving the non-FP8 path that returns a real bgrad.
tests/pytorch/test_numerics.py	Adds focused FP8 delayed-wgrad coverage for consuming the eager bias gradient and clearing the parameter gradient before backward_dw.

Sequence Diagram

sequenceDiagram
  participant Autograd
  participant Linear
  participant Hook as Gradient hook
  participant Store as WeightGradStore
  Autograd->>Linear: backward()
  Linear-->>Autograd: eager FP8 bias gradient
  Autograd->>Hook: accumulate bias.grad
  Hook->>Hook: fold into main_grad and clear bias.grad
  Linear->>Store: backward_dw()
  Store-->>Linear: weight gradient, empty bgrad
  Linear->>Linear: skip empty bias backfill

Loading

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