Implementation:NVIDIA TransformerEngine Ops Fused Backward Activation Bias
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Optimization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Fused backward pass operation that combines activation gradient, bias gradient, and quantization into a single kernel launch for GELU and ReLU activations.
Description
BackwardActivationBias is a FusedOperation that fuses the backward pass of activation (GELU or ReLU) with bias gradient computation and output quantization into a single kernel call. It detects the [quantize_op, Bias, Activation] pattern in the backward operation list and replaces it with the fused tex.dbias_dgelu or tex.dbias_drelu kernel. This reduces memory traffic and kernel launch overhead. The fusion requires that a quantizer is available from the previous operation (Bias context).
Usage
Automatically applied by the operation fuser when it detects a compatible backward pass pattern. Requires a quantization recipe to be active.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/ops/fused/backward_activation_bias.py- Lines
- 1--137
Signature
class BackwardActivationBias(FusedOperation):
def __init__(self, *, bias: Bias, activation: _ActivationOperation): ...
def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs) -> Tuple: ...
@staticmethod
def fuse_backward_ops(ops, *, recipe=None, **unused) -> list[FusibleOperation]: ...
Import
from transformer_engine.pytorch.ops.fused.backward_activation_bias import BackwardActivationBias
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| bias | Bias | Yes | The Bias operation to fuse |
| activation | _ActivationOperation | Yes | GELU or ReLU activation operation |
| grad_output | torch.Tensor | Yes | Upstream gradient |
Outputs
| Name | Type | Description |
|---|---|---|
| grad_input | torch.Tensor | Gradient w.r.t. activation input (may be quantized) |
| grad_bias | torch.Tensor | Gradient w.r.t. bias parameter |
Usage Examples
# Automatically fused by the operation fuser when detecting pattern:
# [op_with_quantizer, Bias, GELU] in the backward pass
# No manual usage required - the fuser handles fusion automatically