Implementation:NVIDIA TransformerEngine JAX Flax Transformer
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Attention |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements complete Flax transformer layer modules including multi-head attention, dot-product attention, and full transformer layers with FP8 support, distributed sharding, and various attention features.
Description
DotProductAttention dispatches between _FusedDotProductAttention (cuDNN-backed) and _UnfusedDotProductAttention (pure JAX) based on kernel availability. MultiHeadAttention handles QKV projection, head splitting, and attention output projection with logical axis sharding. TransformerLayer composes self-attention, optional cross-attention, and MLP blocks with residual connections, drop path, and layer normalization. extend_logical_axis_rules maps TE's logical axes to mesh axes for Flax's partitioning system. Supports rotary position embeddings, relative position biases, sliding window attention, and context parallelism.
This is the highest-level user-facing module, providing complete transformer layer implementations that users can directly stack to build transformer models with FP8 acceleration and distributed training support.
Usage
Use these modules to build complete transformer models in Flax. TransformerLayer provides a complete pre-norm or post-norm transformer layer. MultiHeadAttention can be used standalone for custom attention patterns. Use extend_logical_axis_rules when setting up distributed training with Flax's partitioning.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/flax/transformer.py- Lines
- 1--2319
Signature
def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules: ...
def rotary_pos_emb(x: jnp.ndarray, position: jnp.ndarray) -> jnp.ndarray: ...
class DotProductAttention(nn.Module):
head_dim: int = ...
num_attention_heads: int = ...
num_gqa_groups: int = ...
attn_mask_type: str = "causal"
attn_bias_type: Optional[str] = None
...
class MultiHeadAttention(nn.Module):
head_dim: int = ...
num_attention_heads: int = ...
num_gqa_groups: Optional[int] = None
...
class RelativePositionBiases(nn.Module):
num_buckets: int = 32
max_distance: int = 128
num_attention_heads: int = ...
...
class TransformerLayerType(Enum):
ENCODER = "encoder"
DECODER = "decoder"
class TransformerLayer(nn.Module):
hidden_size: int = 512
mlp_hidden_size: int = 2048
num_attention_heads: int = 8
layernorm_type: str = "layernorm"
hidden_dropout: float = 0.1
attention_dropout: float = 0.1
layer_type: TransformerLayerType = TransformerLayerType.ENCODER
...
Import
from transformer_engine.jax.flax.transformer import TransformerLayer, MultiHeadAttention, DotProductAttention, extend_logical_axis_rules
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| inputs | jnp.ndarray |
Yes | Input tensor of shape [batch, seqlen, hidden] |
| attention_mask | jnp.ndarray |
No | Attention mask tensor |
| encoder_output | jnp.ndarray |
No | Encoder output for cross-attention (decoder layers) |
| deterministic | bool |
No | Whether to disable dropout (default False) |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Transformer layer output of shape [batch, seqlen, hidden] |
Usage Examples
from transformer_engine.jax.flax.transformer import (
TransformerLayer, TransformerLayerType, extend_logical_axis_rules
)
import flax.linen as nn
import jax
# Build a transformer encoder layer
layer = TransformerLayer(
hidden_size=512,
mlp_hidden_size=2048,
num_attention_heads=8,
layernorm_type="rmsnorm",
hidden_dropout=0.1,
attention_dropout=0.1,
layer_type=TransformerLayerType.ENCODER,
mlp_activations=("gelu",),
)
# Extend Flax axis rules for distributed training
rules = extend_logical_axis_rules(flax.linen.partitioning.default_mesh_rules)
# Initialize and run
params = layer.init(rng, input_tensor)
output = layer.apply(params, input_tensor, deterministic=False)