Implementation:NVIDIA TransformerEngine Ops MakeExtraOutput
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Optimization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Fusible operation that exposes an intermediate tensor as an extra output from the operation fuser, enabling residual connections and the BackwardLinearAdd fusion.
Description
MakeExtraOutput is a BasicOperation that returns the input tensor unchanged but also exposes it as an extra output from the operation fuser. In the backward pass, the gradient with respect to the extra output is accumulated with the upstream gradient. When in_place=True, gradient accumulation is performed in-place, which is required for enabling the BackwardLinearAdd fusion pattern. This is the backward-pass counterpart to AddExtraInput.
Usage
Used within the operation fuser to create branch points for residual connections. The in-place mode is an advanced feature for enabling specific fused backward patterns.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/ops/basic/make_extra_output.py- Lines
- 1--95
Signature
class MakeExtraOutput(BasicOperation):
num_extra_outputs: int = 1
def __init__(self, *, in_place: bool = False): ...
def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, prev_op_grad_output_quantizer, next_op_input_quantizer, basic_op_kwargs) -> Tuple[torch.Tensor, Iterable]: ...
def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs) -> Tuple[torch.Tensor, Iterable, Iterable]: ...
Import
from transformer_engine.pytorch.ops.basic.make_extra_output import MakeExtraOutput
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| input_ | torch.Tensor | Yes | Input tensor (passed through as both main and extra output) |
| in_place | bool | No | Whether to accumulate gradients in-place |
Outputs
| Name | Type | Description |
|---|---|---|
| output | torch.Tensor | Same as input (identity) |
| extra_output | torch.Tensor | Copy of the input exposed as extra output |
Usage Examples
from transformer_engine.pytorch.ops.basic.make_extra_output import MakeExtraOutput
from transformer_engine.pytorch.ops import Sequential
model = Sequential(
MakeExtraOutput(in_place=True),
linear_op,
)
output, residual = model(input_tensor)