diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index b079d48aff..60d0536033 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -23,6 +23,7 @@ ) from utils import pytest_parametrize_wrapper from transformer_engine_jax import get_cudnn_version, get_device_compute_capability +from transformer_engine.jax.sharding import MeshResource from transformer_engine.jax.attention import ( is_fused_attn_kernel_available, AttnBiasType, @@ -39,6 +40,25 @@ DTYPES = [jnp.bfloat16] +AxisType = getattr(jax.sharding, "AxisType", None) +EXPLICIT_SHARDING_TEST = pytest.mark.skipif( + AxisType is None or not hasattr(jax, "set_mesh") or not hasattr(jax, "shard_map"), + reason="JAX explicit sharding is unavailable", +) +SHARDING_MODES = [ + pytest.param(False, id="SHARDY"), + pytest.param(True, marks=EXPLICIT_SHARDING_TEST, id="EXPLICIT_SHARDING"), +] +EXPLICIT_SHARDING_MULTI_AXIS_CONFIGS = [ + pytest.param( + 4, + (2, 2), + ("dp", "tpsp"), + MeshResource(dp_resource="dp", tpsp_resource="tpsp"), + id="n4_dp2_tp2", + ) +] + DISTRIBUTED_SELF_ATTN_DATA_SHAPES = { "L0": [()], "L1": [(32, 1024, 16, 128)], @@ -77,9 +97,11 @@ def impl_test_self_attn( attn_mask_type, dtype, softmax_type, + mesh_axis_types=None, + test_backward=True, + is_training=True, ): dropout_prob = 0.0 - is_training = True batch, seqlen, num_head, hidden = data_shape if not is_fused_attn_kernel_available( @@ -131,11 +153,16 @@ def impl_test_self_attn( mesh_shape=mesh_shape, mesh_axes=mesh_axes, mesh_resource=mesh_resource, - coll_count_ref=col_ref, + coll_count_ref=col_ref if test_backward else None, + mesh_axis_types=mesh_axis_types, ) - runner.test_backward() + if test_backward: + runner.test_backward() + else: + runner.test_forward() @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_SELF_ATTN_DATA_SHAPES) @pytest.mark.parametrize( "attn_bias_type, bias_shape", @@ -173,6 +200,7 @@ def test_self_attn( attn_mask_type, dtype, softmax_type, + explicit_sharding, ): self.impl_test_self_attn( device_count, @@ -185,6 +213,33 @@ def test_self_attn( attn_mask_type, dtype, softmax_type, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, + ) + + @EXPLICIT_SHARDING_TEST + @pytest.mark.parametrize( + "device_count,mesh_shape,mesh_axes,mesh_resource", EXPLICIT_SHARDING_MULTI_AXIS_CONFIGS + ) + @pytest.mark.parametrize( + "attn_bias_type", [AttnBiasType.PRE_SCALE_BIAS, AttnBiasType.POST_SCALE_BIAS] + ) + def test_self_attn_explicit_sharding_bias( + self, device_count, mesh_shape, mesh_axes, mesh_resource, attn_bias_type + ): + self.impl_test_self_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + (4, 128, 8, 64), + attn_bias_type, + BiasShape._1HSS, + AttnMaskType.CAUSAL_MASK, + jnp.bfloat16, + AttnSoftmaxType.VANILLA_SOFTMAX, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes), + test_backward=False, + is_training=False, ) @@ -202,21 +257,7 @@ def generate_collectives_count_ref(self): all_reduce_loss_bytes = 4 # 1 * FP32 return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0) - @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) - @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_CROSS_ATTN_DATA_SHAPES) - @pytest.mark.parametrize( - "attn_mask_type", [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK] - ) - @pytest.mark.parametrize("dtype", DTYPES) - @pytest.mark.parametrize( - "softmax_type", - [ - pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), - pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), - pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), - ], - ) - def test_cross_attn( + def impl_test_cross_attn( self, device_count, mesh_shape, @@ -226,6 +267,7 @@ def test_cross_attn( attn_mask_type, dtype, softmax_type, + mesh_axis_types=None, ): attn_bias_type = AttnBiasType.NO_BIAS bias_shape = None @@ -277,9 +319,49 @@ def test_cross_attn( mesh_axes=mesh_axes, mesh_resource=mesh_resource, coll_count_ref=col_ref, + mesh_axis_types=mesh_axis_types, ) runner.test_backward() + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) + @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_CROSS_ATTN_DATA_SHAPES) + @pytest.mark.parametrize( + "attn_mask_type", [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK] + ) + @pytest.mark.parametrize("dtype", DTYPES) + @pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), + pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), + ], + ) + def test_cross_attn( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + attn_mask_type, + dtype, + softmax_type, + explicit_sharding, + ): + self.impl_test_cross_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + attn_mask_type, + dtype, + softmax_type, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, + ) + DISTRIBUTED_SCORE_MOD_DATA_SHAPES = { "L0": [], @@ -431,6 +513,7 @@ def impl_test_context_parallel_attn( num_segments_per_seq=None, return_max_logit=False, check_forward_output=True, + mesh_axis_types=None, ): if qkv_layout.is_thd(): if not load_balanced and ( @@ -486,6 +569,7 @@ def impl_test_context_parallel_attn( mesh_resource=mesh_resource, cp_strategy=cp_strategy, cp_load_balanced=load_balanced, + mesh_axis_types=mesh_axis_types, ) def check_has_backend_for_mask(mask_type): @@ -534,10 +618,74 @@ def check_has_backend_for_mask(mask_type): runner.test_backward() del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] + @EXPLICIT_SHARDING_TEST + @pytest.mark.parametrize("cp_strategy", [CPStrategy.ALL_GATHER, CPStrategy.RING]) + def test_context_parallel_explicit_sharding(self, cp_strategy): + self.impl_test_context_parallel_attn( + 2, + (1, 2, 1), + ("dp", "cp", "tpsp"), + MeshResource(dp_resource="dp", cp_resource="cp", tpsp_resource="tpsp"), + (2, 128, 8, 64), + 2, + AttnMaskType.CAUSAL_MASK, + jnp.bfloat16, + QKVLayout.BSHD_BS2HD, + True, + cp_strategy, + mesh_axis_types=(AxisType.Explicit,) * 3, + ) + + @EXPLICIT_SHARDING_TEST + @pytest.mark.parametrize( + "cp_strategy,window_size,stripe_size", + [ + pytest.param(CPStrategy.ALL_GATHER, (20, 0), 64, id="all-gather-swa"), + pytest.param(CPStrategy.RING, (-1, -1), 1, id="ring"), + ], + ) + def test_context_parallel_thd_explicit_sharding(self, cp_strategy, window_size, stripe_size): + self.impl_test_context_parallel_attn( + 2, + (1, 2, 1), + ("dp", "cp", "tpsp"), + MeshResource(dp_resource="dp", cp_resource="cp", tpsp_resource="tpsp"), + (2, 128, 8, 64), + 1, + AttnMaskType.PADDING_CAUSAL_MASK, + jnp.bfloat16, + QKVLayout.THD_THD_THD, + True, + cp_strategy, + window_size=window_size, + stripe_size=stripe_size, + num_segments_per_seq=5, + mesh_axis_types=(AxisType.Explicit,) * 3, + ) + + @EXPLICIT_SHARDING_TEST + def test_context_parallel_max_logit_explicit_sharding(self): + self.impl_test_context_parallel_attn( + 2, + (1, 2, 1), + ("dp", "cp", "tpsp"), + MeshResource(dp_resource="dp", cp_resource="cp", tpsp_resource="tpsp"), + (2, 128, 8, 64), + 1, + AttnMaskType.CAUSAL_MASK, + jnp.bfloat16, + QKVLayout.BSHD_BSHD_BSHD, + True, + CPStrategy.ALL_GATHER, + return_max_logit=True, + mesh_axis_types=(AxisType.Explicit,) * 3, + ) + @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) @pytest.mark.parametrize("kv_groups", [1, 8]) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @@ -570,6 +718,7 @@ def test_context_parallel_return_max_logit( cp_strategy, window_size, use_scan_ring, + explicit_sharding, ): """Check CP fused attention returns global per-head max_logit.""" is_thd = qkv_layout.is_thd() @@ -606,12 +755,14 @@ def test_context_parallel_return_max_logit( num_segments_per_seq=num_segments_per_seq, return_max_logit=True, check_forward_output=check_forward_output, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, ) @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) @pytest.mark.parametrize("kv_groups", [1, 8]) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @@ -653,6 +804,7 @@ def test_context_parallel_allgather_striped_attn( window_size, stripe_size, num_segments_per_seq, + explicit_sharding, ): if not qkv_layout.is_thd(): pytest.skip("Only THD layout is supported for CP + AG + Striped attention") @@ -671,12 +823,14 @@ def test_context_parallel_allgather_striped_attn( window_size=window_size, stripe_size=stripe_size, num_segments_per_seq=num_segments_per_seq, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, ) @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES) @pytest.mark.parametrize("kv_groups", [1, 8]) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @@ -700,6 +854,7 @@ def test_context_parallel_allgather_attn( dtype, qkv_layout, load_balanced, + explicit_sharding, ): if qkv_layout.is_thd(): pytest.skip("Only BSHD layout is supported for CP + AG + Dual chunk attention") @@ -715,12 +870,14 @@ def test_context_parallel_allgather_attn( qkv_layout, load_balanced, CPStrategy.ALL_GATHER, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, ) @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES) @pytest.mark.parametrize("kv_groups", [1, 8]) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @@ -757,6 +914,7 @@ def test_context_parallel_ring_attn( load_balanced, use_scan, window_size, + explicit_sharding, ): if window_size != (-1, -1) and not qkv_layout.is_thd(): pytest.skip("Sliding window attention is only supported for THD layout") @@ -782,6 +940,7 @@ def test_context_parallel_ring_attn( use_scan_ring=use_scan, window_size=window_size, stripe_size=stripe_size, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, ) # CP ring and all-gather tests for D=256 @@ -805,6 +964,7 @@ def skip_if_d256_cp_unsupported(qkv_layout): "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest_parametrize_wrapper( "data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES, @@ -828,6 +988,7 @@ def test_context_parallel_ring_attn_d256( qkv_layout, attn_mask_type, window_size, + explicit_sharding, ): """D=256 CP ring coverage.""" self.skip_if_d256_cp_unsupported(qkv_layout) @@ -847,12 +1008,14 @@ def test_context_parallel_ring_attn_d256( use_scan_ring=False, window_size=window_size, stripe_size=1 if qkv_layout.is_thd() else None, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, ) @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) + @pytest.mark.parametrize("explicit_sharding", SHARDING_MODES) @pytest_parametrize_wrapper( "data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES, @@ -876,6 +1039,7 @@ def test_context_parallel_allgather_attn_d256( qkv_layout, attn_mask_type, window_size, + explicit_sharding, ): """D=256 CP all-gather coverage.""" self.skip_if_d256_cp_unsupported(qkv_layout) @@ -895,6 +1059,7 @@ def test_context_parallel_allgather_attn_d256( window_size=window_size, stripe_size=128 if qkv_layout.is_thd() else None, num_segments_per_seq=5 if qkv_layout.is_thd() else None, + mesh_axis_types=(AxisType.Explicit,) * len(mesh_axes) if explicit_sharding else None, ) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index d68e409331..f09d6d1bf8 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Tests for fused attention""" import os +from contextlib import contextmanager from enum import Enum, auto from dataclasses import dataclass, field from functools import partial @@ -431,6 +432,7 @@ class FusedAttnRunner: mesh_shape: tuple[int, ...] = (1, 1, 1) mesh_axes: tuple[str, ...] = ("dp", "cp", "tp") mesh_resource: MeshResource = field(default_factory=partial(MeshResource, "dp", "cp", "tp")) + mesh_axis_types: Optional[tuple[Any, ...]] = None # Context parallel aux arguments cp_strategy: CPStrategy = CPStrategy.DEFAULT @@ -623,7 +625,7 @@ def _setup_inputs(self): # Create a mesh for distributed tests self.devices = np.asarray(jax.devices()[: self.number_of_devices]).reshape(*self.mesh_shape) - self.mesh = Mesh(self.devices, self.mesh_axes) + self.mesh = Mesh(self.devices, self.mesh_axes, axis_types=self.mesh_axis_types) self.dp_size = self.mesh.shape.get(self.mesh_resource.dp_resource, 1) self.cp_size = self.mesh.shape.get(self.mesh_resource.cp_resource, 1) self.tp_size = self.mesh.shape.get(self.mesh_resource.tpsp_resource, 1) @@ -949,6 +951,17 @@ def to_dp_shardings(x): self.seq_length_offset_pspec = PartitionSpec(self.mesh_resource.dp_resource, None) self.seq_length_offset_sharding = NamedSharding(self.mesh, self.seq_length_offset_pspec) + @contextmanager + def _mesh_context(self): + """Enter the appropriate mesh context for Auto or explicit axes.""" + mesh_context = self.mesh if self.mesh_axis_types is None else jax.set_mesh(self.mesh) + with mesh_context, autocast(mesh_resource=self.mesh_resource): + yield + + def _assert_explicit_spec(self, value, expected_spec): + if self.mesh_axis_types is not None: + assert jax.typeof(value).sharding.spec == expected_spec + def test_forward(self, return_max_logit=False, check_output=True): """ Test forward with JITted primitive and unJITted reference @@ -1015,10 +1028,15 @@ def test_forward(self, return_max_logit=False, check_output=True): ], ) - with self.mesh, autocast(mesh_resource=self.mesh_resource): + with self._mesh_context(): primitive_out = customcall_fused_dpa_jit(*customcall_args) if return_max_logit: primitive_out, primitive_max_logit = primitive_out + self._assert_explicit_spec( + primitive_max_logit, + PartitionSpec(self.qkvo_psec[-2]), + ) + self._assert_explicit_spec(primitive_out, self.qkvo_psec) primitive_out = self.cp_inverse_reorder_fn(primitive_out) if return_max_logit: @@ -1050,10 +1068,8 @@ def test_forward(self, return_max_logit=False, check_output=True): assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype) if self.coll_count_ref is not None: - with self.mesh, autocast(mesh_resource=self.mesh_resource): - target_hlo = ( - customcall_fused_dpa_jit.lower(*customcall_args, **kwargs).compile().as_text() - ) + with self._mesh_context(): + target_hlo = customcall_fused_dpa_jit.lower(*customcall_args).compile().as_text() assert_equal_collectives(target_hlo, self.coll_count_ref) def test_backward(self, return_max_logit=False): @@ -1196,9 +1212,12 @@ def grad_func( ) ) - with self.mesh, autocast(mesh_resource=self.mesh_resource): + with self._mesh_context(): primitive_out, primitive_dgrad = jitted_primitive(*customcall_args) + for primitive_grad in primitive_dgrad[:3]: + self._assert_explicit_spec(primitive_grad, self.qkvo_psec) + reference_out, reference_dgrad = jitted_reference(*args) # Skip elementwise comparison when dropout enabled @@ -1221,10 +1240,10 @@ def check_dqkv(primitive, reference, pad, idx): _split_valid_and_invalid(primitive, reference, pad) ) - print_debug_tensor_stats(f"primitive_grad_valid[{idx}]", primitive_valid[idx]) - print_debug_tensor_stats(f"reference_grad_valid[{idx}]", reference_valid[idx]) + print_debug_tensor_stats(f"primitive_grad_valid[{idx}]", primitive_valid) + print_debug_tensor_stats(f"reference_grad_valid[{idx}]", reference_valid) print_debug_tensor_stats( - f"diff_grad[{idx}]", jnp.abs(primitive_valid[idx] - reference_valid[idx]) + f"diff_grad[{idx}]", jnp.abs(primitive_valid - reference_valid) ) assert_allclose( @@ -1297,7 +1316,7 @@ def check_dqkv(primitive, reference, pad, idx): ) if self.coll_count_ref is not None: - with self.mesh, autocast(mesh_resource=self.mesh_resource): + with self._mesh_context(): target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text() assert_equal_collectives(target_hlo, self.coll_count_ref) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 6d7c823e12..ae0ad62985 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -387,6 +387,17 @@ def _obtain_batch_and_max_seqlen(qkv, qkv_layout): return batch, q_max_seqlen, kv_max_seqlen +def _reorder_with_explicit_sharding(op, tensor): + """Run a reorder with automatic axes while preserving typed sharding.""" + sharding = getattr(jax.typeof(tensor), "sharding", None) + mesh = getattr(sharding, "mesh", None) + axis_types = getattr(mesh, "axis_types", ()) + if axis_types and any(axis_type.name == "Explicit" for axis_type in axis_types): + with jax.sharding.use_abstract_mesh(mesh): + return jax.sharding.auto_axes(op, out_sharding=sharding.spec)(tensor) + return op(tensor) + + def reorder_causal_load_balancing( tensor, strategy: ReorderStrategy, cp_size: int, seq_dim: int, stripe_size: int | None = None ): @@ -397,7 +408,13 @@ def reorder_causal_load_balancing( f"Incorrect value for CP dual chunk reordering {stripe_size=}. stripe_size must be" " None" ) - return tex.attention.reorder_causal_dual_chunk_swap(tensor, cp_size, seq_dim, False) + op = partial( + tex.attention.reorder_causal_dual_chunk_swap, + cp_size=cp_size, + seq_dim=seq_dim, + to_contiguous=False, + ) + return _reorder_with_explicit_sharding(op, tensor) if strategy == ReorderStrategy.Striped: # stripe_size > 1 is only supported for CP+THD+AG+Striped>1+SWA # stripe_size = 128 is recommended for CP+THD+AG+Striped>1+SWA @@ -408,9 +425,14 @@ def reorder_causal_load_balancing( ) # Supporting old API defaults of stripe_size=1 effective_stripe_size = 1 if stripe_size is None else stripe_size - return tex.attention.reorder_causal_striped( - tensor, cp_size, seq_dim, False, effective_stripe_size + op = partial( + tex.attention.reorder_causal_striped, + cp_size=cp_size, + seq_dim=seq_dim, + is_inverse=False, + stripe_size=effective_stripe_size, ) + return _reorder_with_explicit_sharding(op, tensor) raise ValueError(f"Unsupported {strategy=}") @@ -424,7 +446,13 @@ def inverse_reorder_causal_load_balancing( f"Incorrect value for CP dual chunk reordering {stripe_size=}. stripe_size must be" " None" ) - return tex.attention.reorder_causal_dual_chunk_swap(tensor, cp_size, seq_dim, True) + op = partial( + tex.attention.reorder_causal_dual_chunk_swap, + cp_size=cp_size, + seq_dim=seq_dim, + to_contiguous=True, + ) + return _reorder_with_explicit_sharding(op, tensor) if strategy == ReorderStrategy.Striped: # stripe_size > 1 is only supported for CP+THD+AG+Striped>1+SWA # stripe_size = 128 is recommended for CP+THD+AG+Striped>1+SWA @@ -435,9 +463,14 @@ def inverse_reorder_causal_load_balancing( ) # Supporting old API defaults of stripe_size=1 effective_stripe_size = 1 if stripe_size is None else stripe_size - return tex.attention.reorder_causal_striped( - tensor, cp_size, seq_dim, True, effective_stripe_size + op = partial( + tex.attention.reorder_causal_striped, + cp_size=cp_size, + seq_dim=seq_dim, + is_inverse=True, + stripe_size=effective_stripe_size, ) + return _reorder_with_explicit_sharding(op, tensor) raise ValueError(f"Unsupported {strategy=}") diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 6ea54195ab..b96e878213 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -7,6 +7,7 @@ import warnings from dataclasses import dataclass, replace from functools import partial, reduce +from types import SimpleNamespace from typing import Optional, Tuple import jax @@ -43,7 +44,6 @@ all_reduce_sum_along_dp_fsdp, get_mesh_axis_size, get_mesh_axis_rank, - get_mesh_axis_rank_host, get_all_mesh_axes, num_of_devices, with_sharding_constraint, @@ -103,6 +103,93 @@ class _FusedAttnConfig: return_max_logit: bool = False +def _explicit_aval_spec(aval): + """Return a rank-padded spec when ``aval`` uses sharding-in-types.""" + sharding = getattr(aval, "sharding", None) + mesh = getattr(sharding, "mesh", None) + axis_types = getattr(mesh, "axis_types", ()) + if not axis_types or not any(axis_type.name == "Explicit" for axis_type in axis_types): + return None + spec = tuple(sharding.spec) + return spec + (None,) * (aval.ndim - len(spec)) + + +def _update_aval_with_spec(aval, *, shape, dtype, spec=None): + """Update an aval and replace typed sharding when its rank/layout changes.""" + kwargs = {"shape": shape, "dtype": dtype} + if _explicit_aval_spec(aval) is not None: + if spec is None: + spec = (None,) * len(shape) + kwargs["sharding"] = NamedSharding(aval.sharding.mesh, PartitionSpec(*spec)) + return aval.update(**kwargs) + + +def _explicit_value_pspec(value): + """Return a PartitionSpec for an explicitly sharded value, else ``None``.""" + spec = _explicit_aval_spec(jax.typeof(value)) + return None if spec is None else PartitionSpec(*spec) + + +def _fused_attn_fwd_explicit_out_specs(q, config): + """Output specs for the explicit fused-attention boundary.""" + q_spec = _explicit_aval_spec(jax.typeof(q)) + if q_spec is None: + return None + + output_spec = (*q_spec[:-3], *q_spec[-2:]) if config.qkv_layout.is_qkvpacked() else q_spec + is_packed_softmax = get_cudnn_version() >= (9, 6, 0) and config.qkv_layout.is_thd() + if config.qkv_layout.is_qkvpacked(): + if is_packed_softmax: + softmax_spec = (*q_spec[:-4], q_spec[-4], q_spec[-2], None) + else: + softmax_spec = (*q_spec[:-4], q_spec[-2], q_spec[-4], None) + elif is_packed_softmax: + softmax_spec = (*q_spec[:-3], q_spec[-3], q_spec[-2], None) + else: + softmax_spec = (*q_spec[:-3], q_spec[-2], q_spec[-3], None) + + rng_spec = (tuple(jax.typeof(q).sharding.mesh.axis_names), None) + max_logit_spec = (output_spec[-2],) if config.return_max_logit else (None,) + return [PartitionSpec(*spec) for spec in (output_spec, softmax_spec, rng_spec, max_logit_spec)] + + +def _run_explicit_partitioned(primitive_cls, config, args): + """Run an existing attention partition implementation via ``shard_map``.""" + arg_avals = tuple(jax.typeof(arg) for arg in args) + mesh = arg_avals[0].sharding.mesh + out_avals = primitive_cls.outer_abstract(*arg_avals, config=config) + + def to_info(aval): + return SimpleNamespace( + sharding=aval.sharding, + shape=aval.shape, + ndim=aval.ndim, + dtype=aval.dtype, + ) + + arg_infos = tuple(to_info(aval) for aval in arg_avals) + result_infos = tuple(to_info(aval) for aval in out_avals) + _, impl, out_shardings, arg_shardings = primitive_cls.partition( + config, mesh, arg_infos, result_infos + ) + in_specs = tuple( + PartitionSpec(*tuple(sharding.spec)[: aval.ndim]) + for sharding, aval in zip(arg_shardings, arg_avals) + ) + out_specs = tuple( + PartitionSpec(*tuple(sharding.spec)[: aval.ndim]) + for sharding, aval in zip(out_shardings, out_avals) + ) + args = tuple(jax.sharding.reshard(arg, spec) for arg, spec in zip(args, in_specs)) + return jax.shard_map( + impl, + mesh=mesh, + in_specs=in_specs, + out_specs=out_specs, + check_vma=False, + )(*args) + + @dataclass(frozen=True) class FusedAttnHelper: """ @@ -336,7 +423,14 @@ def abstract( ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) output_shape = (*batch_shape, q_max_seqlen, attn_heads, v_head_dim) - out_aval = q_aval.update(shape=output_shape, dtype=q_dtype) + q_spec = _explicit_aval_spec(q_aval) + if q_spec is not None and config.qkv_layout.is_qkvpacked(): + output_spec = (*q_spec[:-3], *q_spec[-2:]) + else: + output_spec = q_spec + out_aval = _update_aval_with_spec( + q_aval, shape=output_shape, dtype=q_dtype, spec=output_spec + ) # backend determines the softmax buffer shape/dtype backend = FusedAttnHelper( @@ -358,6 +452,7 @@ def abstract( config.return_max_logit, ).get_fused_attn_backend() + is_packed_softmax = get_cudnn_version() >= (9, 6, 0) and config.qkv_layout.is_thd() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: # cuDNN 9.6 reduces the required softmax shape if get_cudnn_version() >= (9, 6, 0): @@ -375,7 +470,20 @@ def abstract( softmax_dtype = dtypes.canonicalize_dtype(jnp.float32) else: raise ValueError(f"Unsupported {backend=}") - softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype) + if q_spec is None: + softmax_spec = None + elif config.qkv_layout.is_qkvpacked(): + if is_packed_softmax: + softmax_spec = (*q_spec[:-4], q_spec[-4], q_spec[-2], None) + else: + softmax_spec = (*q_spec[:-4], q_spec[-2], q_spec[-4], None) + elif is_packed_softmax: + softmax_spec = (*q_spec[:-3], q_spec[-3], q_spec[-2], None) + else: + softmax_spec = (*q_spec[:-3], q_spec[-2], q_spec[-3], None) + softmax_aux_aval = _update_aval_with_spec( + q_aval, shape=softmax_shape, dtype=softmax_dtype, spec=softmax_spec + ) if config.return_max_logit: # cuDNN Max is row-wise over S_kv. Dense and SM120 THD use # [..., H, S_q, 1]; cuDNN >= 9.6 non-SM120 THD uses [..., S_q, H, 1]. @@ -386,7 +494,10 @@ def abstract( max_tensor_shape = (*batch_shape, attn_heads, q_max_seqlen, 1) else: max_tensor_shape = (0,) - max_tensor_aval = q_aval.update(shape=max_tensor_shape, dtype=softmax_dtype) + max_tensor_spec = softmax_spec if config.return_max_logit else (None,) + max_tensor_aval = _update_aval_with_spec( + q_aval, shape=max_tensor_shape, dtype=softmax_dtype, spec=max_tensor_spec + ) # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with # 32-bit unsigned int to get the buffer size we need in the C++ kernel @@ -396,7 +507,16 @@ def abstract( seed_dtype == checker.rng_state_dtype ), f"Expected seed_dtype={checker.rng_state_dtype}, but got seed_dtype={seed_dtype}" rng_state_shape = (seed_aval.shape[0], checker.rng_state_size) - rng_state_aval = seed_aval.update(shape=rng_state_shape, dtype=checker.rng_state_dtype) + seed_spec = _explicit_aval_spec(seed_aval) + rng_state_spec = ( + None if seed_spec is None else (tuple(seed_aval.sharding.mesh.axis_names), None) + ) + rng_state_aval = _update_aval_with_spec( + seed_aval, + shape=rng_state_shape, + dtype=checker.rng_state_dtype, + spec=rng_state_spec, + ) if config.attn_bias_type == AttnBiasType.NO_BIAS: bias_batch = bias_heads = 0 @@ -436,8 +556,11 @@ def abstract( config.return_max_logit, bottom_right_diagonal, ) - wkspace_aval = q_aval.update( - shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) + wkspace_aval = _update_aval_with_spec( + q_aval, + shape=wkspace_info[0], + dtype=te_dtype_to_jax_dtype(wkspace_info[1]), + spec=(None,) * len(wkspace_info[0]), ) assert ( @@ -465,7 +588,13 @@ def outer_abstract(*args, **kwargs): *args, **kwargs ) max_logit_shape = (out_aval.shape[-2],) if kwargs["config"].return_max_logit else (0,) - max_logit_aval = out_aval.update(shape=max_logit_shape, dtype=out_aval.dtype) + out_spec = _explicit_aval_spec(out_aval) + max_logit_spec = None + if out_spec is not None: + max_logit_spec = (out_spec[-2],) if kwargs["config"].return_max_logit else (None,) + max_logit_aval = _update_aval_with_spec( + out_aval, shape=max_logit_shape, dtype=out_aval.dtype, spec=max_logit_spec + ) return out_aval, softmax_aux_aval, rng_state_aval, max_logit_aval @staticmethod @@ -990,12 +1119,26 @@ def abstract( config.bottom_right_diagonal, ) - dq_aval = q_aval.update(shape=q_aval.shape, dtype=q_dtype) - dk_aval = k_aval.update(shape=k_aval.shape, dtype=k_dtype) - dv_aval = v_aval.update(shape=v_aval.shape, dtype=v_dtype) - dbias_aval = bias_aval.update(shape=bias_aval.shape, dtype=bias_dtype) - wkspace_aval = q_aval.update( - shape=wkspace_shape, dtype=te_dtype_to_jax_dtype(wkspace_dtype) + dq_aval = _update_aval_with_spec( + q_aval, shape=q_aval.shape, dtype=q_dtype, spec=_explicit_aval_spec(q_aval) + ) + dk_aval = _update_aval_with_spec( + k_aval, shape=k_aval.shape, dtype=k_dtype, spec=_explicit_aval_spec(k_aval) + ) + dv_aval = _update_aval_with_spec( + v_aval, shape=v_aval.shape, dtype=v_dtype, spec=_explicit_aval_spec(v_aval) + ) + dbias_aval = _update_aval_with_spec( + bias_aval, + shape=bias_aval.shape, + dtype=bias_dtype, + spec=_explicit_aval_spec(bias_aval), + ) + wkspace_aval = _update_aval_with_spec( + q_aval, + shape=wkspace_shape, + dtype=te_dtype_to_jax_dtype(wkspace_dtype), + spec=(None,) * len(wkspace_shape), ) # Validate incoming softmax_offset shape and dtype @@ -1014,11 +1157,19 @@ def abstract( ) if config.softmax_type == AttnSoftmaxType.VANILLA_SOFTMAX: - dsoftmax_offset_aval = q_aval.update( - shape=softmax_offset_aval.shape, dtype=softmax_offset_aval.dtype + dsoftmax_offset_aval = _update_aval_with_spec( + softmax_offset_aval, + shape=softmax_offset_aval.shape, + dtype=softmax_offset_aval.dtype, + spec=_explicit_aval_spec(softmax_offset_aval), ) else: - dsoftmax_offset_aval = q_aval.update(shape=(1, attn_heads, 1, 1), dtype=jnp.float32) + dsoftmax_offset_aval = _update_aval_with_spec( + softmax_offset_aval, + shape=(1, attn_heads, 1, 1), + dtype=jnp.float32, + spec=_explicit_aval_spec(softmax_offset_aval), + ) return dq_aval, dk_aval, dv_aval, dbias_aval, dsoftmax_offset_aval, wkspace_aval @@ -3291,7 +3442,6 @@ def fwd_impl( subblock_config = config cp_size = get_mesh_axis_size(config.cp_axis, mesh) - cp_rank = get_mesh_axis_rank_host(config.cp_axis, mesh) cp_perm = [(i, (i + 1) % cp_size) for i in range(cp_size)] batch, q_max_seqlen, head, _ = q.shape @@ -3333,19 +3483,27 @@ def compute(config): ) if config.window_size != (-1, -1): - kv_src_rank = (cp_size + cp_rank - idx) % cp_size - # Note: all inputs of adjust_cp_striped_window_size should be host values - cp_striped_window_size = adjust_cp_striped_window_size( - cp_rank, kv_src_rank, cp_size, config.window_size - ) - current_config = replace( - subblock_config, cp_striped_window_size=cp_striped_window_size + cp_rank = get_mesh_axis_rank(config.cp_axis, mesh) + rank_configs = [] + for rank in range(cp_size): + kv_src_rank = (cp_size + rank - idx) % cp_size + cp_striped_window_size = adjust_cp_striped_window_size( + rank, kv_src_rank, cp_size, config.window_size + ) + rank_configs.append( + replace( + subblock_config, + cp_striped_window_size=cp_striped_window_size, + ) + ) + output_per_step, softmax_aux_per_step, _, max_logit_per_step = lax.switch( + cp_rank, + tuple(partial(compute, rank_config) for rank_config in rank_configs), ) else: - current_config = subblock_config - output_per_step, softmax_aux_per_step, _, max_logit_per_step = compute( - current_config - ) + output_per_step, softmax_aux_per_step, _, max_logit_per_step = compute( + subblock_config + ) softmax_aux_per_step = softmax_aux_per_step.reshape((batch, q_max_seqlen, head, 1)) @@ -3460,8 +3618,6 @@ def bwd_impl( subblock_config = config cp_size = get_mesh_axis_size(config.cp_axis, mesh) - # We need cp_rank to be a host value for adjust_cp_striped_window_size() - cp_rank = get_mesh_axis_rank_host(config.cp_axis, mesh) cp_perm = [(i, (i + 1) % cp_size) for i in range(cp_size)] dq = jnp.zeros_like(q) @@ -3502,17 +3658,25 @@ def compute(config): return dq_per_step, dkv_per_step, dbias_per_step if config.window_size != (-1, -1): - kv_src_rank = (cp_size + cp_rank - idx) % cp_size - # Note: all inputs of adjust_cp_striped_window_size should be host values - cp_striped_window_size = adjust_cp_striped_window_size( - cp_rank, kv_src_rank, cp_size, config.window_size - ) - current_config = replace( - subblock_config, cp_striped_window_size=cp_striped_window_size + cp_rank = get_mesh_axis_rank(config.cp_axis, mesh) + rank_configs = [] + for rank in range(cp_size): + kv_src_rank = (cp_size + rank - idx) % cp_size + cp_striped_window_size = adjust_cp_striped_window_size( + rank, kv_src_rank, cp_size, config.window_size + ) + rank_configs.append( + replace( + subblock_config, + cp_striped_window_size=cp_striped_window_size, + ) + ) + dq_per_step, dkv_per_step, dbias_per_step = lax.switch( + cp_rank, + tuple(partial(compute, rank_config) for rank_config in rank_configs), ) else: - current_config = subblock_config - dq_per_step, dkv_per_step, dbias_per_step = compute(current_config) + dq_per_step, dkv_per_step, dbias_per_step = compute(subblock_config) kv_next, dkv = jnp.unstack(kv_dkv) dq += dq_per_step @@ -3696,30 +3860,40 @@ def fused_attn_fwd( return_max_logit=return_max_logit, ) - primitive = None + primitive_cls = None match context_parallel_strategy: case CPStrategy.DEFAULT | CPStrategy.ALL_GATHER: if qkv_layout.is_thd(): - primitive = FusedAttnCPStripedWithAllGatherFwdPrimitive.outer_primitive + primitive_cls = FusedAttnCPStripedWithAllGatherFwdPrimitive else: - primitive = FusedAttnCPWithAllGatherFwdPrimitive.outer_primitive + primitive_cls = FusedAttnCPWithAllGatherFwdPrimitive case CPStrategy.RING: # We must use stripe attention for THD-RING if qkv_layout.is_thd(): - primitive = FusedRingAttnStripedFwdPrimitive.outer_primitive + primitive_cls = FusedRingAttnStripedFwdPrimitive else: - primitive = FusedRingAttnFwdPrimitive.outer_primitive + primitive_cls = FusedRingAttnFwdPrimitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - output, softmax_aux, rng_state, max_logit = primitive.bind( + primitive_args = ( *qkv_for_primitive, bias, softmax_offset, seed, *seq_desc_flatten, - config=fused_config, ) - rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None)) + + def bind_primitive(*args): + return primitive_cls.outer_primitive.bind(*args, config=fused_config) + + explicit_out_specs = _fused_attn_fwd_explicit_out_specs(qkv_for_primitive[0], fused_config) + if explicit_out_specs is None: + output, softmax_aux, rng_state, max_logit = bind_primitive(*primitive_args) + rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None)) + else: + output, softmax_aux, rng_state, max_logit = _run_explicit_partitioned( + primitive_cls, fused_config, primitive_args + ) return (output, softmax_aux, rng_state, max_logit) @@ -3871,21 +4045,21 @@ def fused_attn_bwd( stripe_size=stripe_size, ) - primitive = None + primitive_cls = None match context_parallel_strategy: case CPStrategy.DEFAULT | CPStrategy.ALL_GATHER: if qkv_layout.is_thd(): - primitive = FusedAttnCPStripedWithAllGatherBwdPrimitive.outer_primitive + primitive_cls = FusedAttnCPStripedWithAllGatherBwdPrimitive else: - primitive = FusedAttnCPWithAllGatherBwdPrimitive.outer_primitive + primitive_cls = FusedAttnCPWithAllGatherBwdPrimitive case CPStrategy.RING: if qkv_layout.is_thd(): - primitive = FusedRingAttnStripedBwdPrimitive.outer_primitive + primitive_cls = FusedRingAttnStripedBwdPrimitive else: - primitive = FusedRingAttnBwdPrimitive.outer_primitive + primitive_cls = FusedRingAttnBwdPrimitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - *qkv_grads, bias_grad, softmax_offset_grad = primitive.bind( + primitive_args = ( *qkv_for_primitive, bias, softmax_offset, @@ -3894,6 +4068,18 @@ def fused_attn_bwd( output, doutput, *seq_desc_flatten, - config=fused_config, ) + + def bind_primitive(*args): + return primitive_cls.outer_primitive.bind(*args, config=fused_config) + + explicit_out_specs = [ + _explicit_value_pspec(value) for value in (*qkv_for_primitive, bias, softmax_offset) + ] + if any(spec is not None for spec in explicit_out_specs): + *qkv_grads, bias_grad, softmax_offset_grad = _run_explicit_partitioned( + primitive_cls, fused_config, primitive_args + ) + else: + *qkv_grads, bias_grad, softmax_offset_grad = bind_primitive(*primitive_args) return tuple(qkv_grads[: len(qkv)]), bias_grad, softmax_offset_grad diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 2e8e611fa3..5317247497 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -164,6 +164,11 @@ def filter_manual_axes(name_or_tuple): return x cleaned_pspec = PartitionSpec(*cleaned_axis_names) + abstract_mesh = get_abstract_mesh() + if abstract_mesh.axis_types and all( + axis_type.name == "Explicit" for axis_type in abstract_mesh.axis_types + ): + return jax.sharding.reshard(x, cleaned_pspec) return jax.lax.with_sharding_constraint(x, cleaned_pspec)