Implementation:NVIDIA TransformerEngine JAX LayerNorm MLP
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Normalization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements a fused layer normalization followed by a two-layer MLP block (GEMM1 + activation + GEMM2) as a single differentiable JAX function with FP8 quantization, sharding, and checkpointing support.
Description
The public layernorm_mlp() delegates to _layernorm_mlp via jax.custom_vjp with many non-diff arguments. The forward pass performs: (1) normalization via tex.normalization_fwd, (2) first GEMM with bias via tex.gemm, (3) activation via tex.act_lu, (4) second GEMM with bias via tex.gemm. JAX checkpoint names (ffn1_ckpt_name, ffn2_ckpt_name) are applied for gradient checkpointing. The backward pass computes gradients for all parameters using fused backward operations (tex.quantize_dact_dbias) for the activation gradient + bias reduction + quantization.
This is the most complex fused operation in the JAX backend, combining four operations (norm, GEMM, activation, GEMM) into a single differentiable unit. This is the standard MLP pattern in transformers, and the fusion maximizes memory efficiency and GPU utilization.
Usage
Use this function when implementing fused LayerNorm+MLP blocks in transformers. It is the primary building block for the LayerNormMLP Flax module and provides optimal performance for the feedforward sublayer.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/layernorm_mlp.py- Lines
- 1--575
Signature
def layernorm_mlp(
x: jnp.ndarray,
gamma: jnp.ndarray,
beta: jnp.ndarray,
kernels: List[jnp.ndarray],
biases: List[jnp.ndarray],
norm_type: str,
zero_centered_gamma: bool = False,
epsilon: float = 1e-6,
transpose_batch_sequence: bool = False,
norm_input_axes: Tuple[str, ...] = None,
dot_1_input_axes: Tuple[str, ...] = None,
dot_2_input_axes: Tuple[str, ...] = None,
kernel_1_axes: Tuple[str, ...] = None,
kernel_2_axes: Tuple[str, ...] = None,
ffn1_ckpt_name: str = "ffn1",
ffn2_ckpt_name: str = "ffn2",
activation_type: Sequence[Union[str, Callable]] = ("relu",),
quantizer_set: QuantizerSet = noop_quantizer_set,
act_params: Optional[ActivationParams] = None,
) -> jnp.ndarray: ...
Import
from transformer_engine.jax.layernorm_mlp import layernorm_mlp
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| x | jnp.ndarray |
Yes | Input tensor |
| gamma | jnp.ndarray |
Yes | Layer normalization scale parameter |
| beta | jnp.ndarray |
Yes | Layer normalization shift parameter |
| kernels | List[jnp.ndarray] |
Yes | Weight matrices [kernel1, kernel2] for the two GEMM layers |
| biases | List[jnp.ndarray] |
Yes | Bias vectors [bias1, bias2] (can contain None) |
| norm_type | str |
Yes | Normalization type: "layernorm" or "rmsnorm" |
| activation_type | Sequence[Union[str, Callable]] |
No | Activation function(s) between GEMM layers (default ("relu",)) |
| quantizer_set | QuantizerSet |
No | FP8 quantizer set |
| ffn1_ckpt_name | str |
No | Checkpoint name for first FFN layer (default "ffn1") |
| ffn2_ckpt_name | str |
No | Checkpoint name for second FFN layer (default "ffn2") |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Result of LayerNorm(x) -> GEMM1 -> Activation -> GEMM2 |
Usage Examples
from transformer_engine.jax.layernorm_mlp import layernorm_mlp
# Fused LayerNorm + MLP forward pass with GELU activation
output = layernorm_mlp(
x=input_tensor,
gamma=ln_scale,
beta=ln_bias,
kernels=[ffn1_weight, ffn2_weight],
biases=[ffn1_bias, ffn2_bias],
norm_type="layernorm",
activation_type=("gelu",),
quantizer_set=fp8_quantizer_set,
)