Implementation:NVIDIA TransformerEngine JAX Sharding
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Manages tensor sharding and distributed training configuration, providing the mapping between TE's logical axes and physical mesh axes for data parallelism, tensor parallelism, FSDP, and context parallelism.
Description
MeshResource is a dataclass holding mesh axis names for each parallelism dimension (dp, tp, fsdp, pp, cp, tpsp). The module defines logical axis constants (BATCH_AXES, SEQLEN_AXES, HEAD_AXES, HIDDEN_TP_AXES, W_FSDP_AXES, etc.) and get_sharding_map_logic_axis_to_mesh_axis maps them to physical mesh axes based on the current MeshResource configuration. with_sharding_constraint and with_sharding_constraint_by_logical_axes apply JAX sharding constraints while handling manual/abstract mesh axes. Helper functions (lax_paral_op, all_reduce_sum_along_dp_fsdp, etc.) perform collective operations along specific mesh dimensions.
This is core infrastructure for distributed training -- every tensor operation in TE uses these utilities to ensure correct data partitioning across GPUs, supporting data parallelism, tensor parallelism, sequence parallelism, FSDP, pipeline parallelism, and context parallelism.
Usage
Use MeshResource to configure the mapping between parallelism dimensions and physical mesh axes. Use global_shard_guard as a context manager to set the global mesh resource. Use with_sharding_constraint_by_logical_axes to apply sharding constraints to tensors.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/sharding.py- Lines
- 1--443
Signature
class MeshResource:
dp_resource: str = None
tp_resource: str = None
fsdp_resource: str = None
pp_resource: str = None
cp_resource: str = None
tpsp_resource: str = None
def is_mesh_available() -> bool: ...
def get_sharding_map_logic_axis_to_mesh_axis() -> Dict: ...
def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec): ...
def with_sharding_constraint_by_logical_axes(
x: jnp.array, logical_axis_names: Tuple[str, ...],
) -> jnp.array: ...
def get_all_mesh_axes() -> Tuple[str, ...]: ...
def lax_paral_op(x: jnp.array, ops, mesh_axis: str) -> jnp.array: ...
def num_of_devices() -> int: ...
def global_shard_guard(resource: MeshResource) -> contextmanager: ...
def global_mesh_resource() -> MeshResource: ...
def all_reduce_sum_along_dp_fsdp(x: jnp.array, mesh: jax.sharding.Mesh) -> jnp.array: ...
Import
from transformer_engine.jax.sharding import MeshResource, global_shard_guard, with_sharding_constraint_by_logical_axes
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| dp_resource | str |
No | Mesh axis name for data parallelism |
| tp_resource | str |
No | Mesh axis name for tensor parallelism |
| fsdp_resource | str |
No | Mesh axis name for FSDP |
| pp_resource | str |
No | Mesh axis name for pipeline parallelism |
| cp_resource | str |
No | Mesh axis name for context parallelism |
Outputs
| Name | Type | Description |
|---|---|---|
| sharded_tensor | jnp.array |
Tensor with sharding constraints applied |
Usage Examples
from transformer_engine.jax.sharding import MeshResource, global_shard_guard
# Configure mesh resources for 2-way TP + 4-way DP
mesh_resource = MeshResource(
dp_resource="dp",
tp_resource="tp",
fsdp_resource=None,
)
# Set global sharding configuration
with global_shard_guard(mesh_resource):
# All TE operations will use this mesh resource
output = model.apply(params, input_data)