Implementation:NVIDIA TransformerEngine Ops Dropout

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

Overview

Fusible dropout operation that randomly zeros tensor entries during training, with support for fused CUDA kernels and CPU activation offloading.

Description

Dropout is a BasicOperation that implements dropout with automatic selection between three implementations: evaluation mode (no-op), fused CUDA kernel (tex.dropout_fwd/tex.dropout_bwd for aligned FP16/BF16 tensors), and unfused Bernoulli masking. The fused path requires tensor element count divisible by 16 and FP16 or BF16 dtype. The dropout mask is saved for the backward pass and supports CPU offloading when enabled.

Usage

Used within the operation fuser pipeline for attention dropout and hidden dropout in transformer models.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/basic/dropout.py

Lines

1--105

Signature

class Dropout(BasicOperation):
    def __init__(self, p: float) -> 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.dropout import Dropout

I/O Contract

Inputs

Name	Type	Required	Description
p	float	Yes	Dropout probability
input_	torch.Tensor	Yes	Input tensor

Outputs

Name	Type	Description
output	torch.Tensor	Tensor with entries randomly zeroed and scaled by 1/(1-p)

Usage Examples

from transformer_engine.pytorch.ops.basic.dropout import Dropout

dropout_op = Dropout(p=0.1)
# In fuser pipeline during training: applies dropout
# During evaluation: identity operation

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements