Implementation:NVIDIA TransformerEngine JAX Checkpoint Policies
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Defines JAX activation checkpoint policies that recognize Transformer Engine GEMM primitives as saveable operations, extending JAX's built-in dot checkpoint policies.
Description
te_gemms_saveable checks whether a given primitive is a TE GemmPrimitive or GroupedGemmPrimitive (or JAX's scaled_matmul_wrapper). Two composite policies are created using jax.checkpoint_policies.save_from_both_policies: dots_and_te_gemms_with_no_batch_dims (combines JAX's no-batch-dims dot policy with TE GEMM recognition) and checkpoint_dots_and_te_gemms (combines JAX's standard dot policy with TE GEMM recognition).
This enables memory-efficient gradient checkpointing in transformer training by allowing JAX's rematerialization system to correctly identify and save/recompute TE GEMM operations alongside standard JAX dot products.
Usage
Use these checkpoint policies with jax.checkpoint (or flax.linen.remat) when training large transformer models that use Transformer Engine's GEMM primitives. This ensures that TE's custom GEMM operations are treated the same as standard JAX dots for checkpointing purposes.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/checkpoint_policies.py- Lines
- 1--38
Signature
def te_gemms_saveable(prim, *_, **__) -> bool:
"""Checkpoint policy for Transformer Engine GEMMs."""
...
dots_and_te_gemms_with_no_batch_dims = jax.checkpoint_policies.save_from_both_policies(
jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims,
te_gemms_saveable,
)
checkpoint_dots_and_te_gemms = jax.checkpoint_policies.save_from_both_policies(
jax.checkpoint_policies.checkpoint_dots,
te_gemms_saveable,
)
Import
from transformer_engine.jax.checkpoint_policies import checkpoint_dots_and_te_gemms, dots_and_te_gemms_with_no_batch_dims
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| prim | jax.core.Primitive |
Yes | JAX primitive to check |
Outputs
| Name | Type | Description |
|---|---|---|
| is_saveable | bool |
Whether the primitive should be saved (True) or recomputed (False) |
Usage Examples
from transformer_engine.jax.checkpoint_policies import checkpoint_dots_and_te_gemms
import jax
# Use with jax.checkpoint for gradient checkpointing
@jax.checkpoint(policy=checkpoint_dots_and_te_gemms)
def transformer_layer(x, params):
# TE GEMM operations will be saved, other ops recomputed
return layer_fn(x, params)
# Or with Flax remat
import flax.linen as nn
layer = nn.remat(TransformerLayer, policy=checkpoint_dots_and_te_gemms)