Implementation:NVIDIA TransformerEngine Ops Bias

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch
Last Updated	2026-02-07 14:00 GMT

Overview

Fusible operation that applies an additive bias, equivalent to the bias in torch.nn.Linear, with support for tensor parallelism.

Description

Bias is a BasicOperation that adds a learnable bias vector to the input tensor's inner dimension. It supports tensor parallelism by distributing the bias along the inner dimension across processes. The backward pass computes the bias gradient using tex.bgrad_quantize when a quantizer is available, enabling fused bias-gradient-quantize operations. The bias parameter is initialized to zeros and supports deferred initialization with meta device.

Usage

Used within the operation fuser pipeline, typically following a BasicLinear operation. Can be fused into ForwardLinearBiasActivation or ForwardLinearBiasAdd patterns.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/basic/bias.py

Lines

1--144

Signature

class Bias(BasicOperation):
    def __init__(self, size, *, device=None, dtype=None, tensor_parallel=False, tensor_parallel_group=None) -> None: ...
    def reset_parameters(self) -> None: ...
    def op_forward(self, ctx, input_, prev_op_grad_output_quantizer, next_op_input_quantizer) -> torch.Tensor: ...
    def op_backward(self, ctx, grad_output) -> Tuple[torch.Tensor, Tuple]: ...

Import

from transformer_engine.pytorch.ops.basic.bias import Bias

I/O Contract

Inputs

Name	Type	Required	Description
size	int	Yes	Inner dimension of input tensor (bias vector size)
input_	torch.Tensor	Yes	Input tensor to add bias to
tensor_parallel	bool	No	Whether to distribute bias across processes

Outputs

Name	Type	Description
output	torch.Tensor	Input tensor with bias added

Usage Examples

from transformer_engine.pytorch.ops.basic.bias import Bias

bias_op = Bias(size=4096, device="cuda", dtype=torch.bfloat16)
output = bias_op.op_forward(ctx, input_tensor, None, None)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements