diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 0b8bb25f3f..d3effcdc07 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -36,6 +36,7 @@ ep_combine_fwd, get_ep_config, ) +from transformer_engine.jax.version_utils import is_collective_stream_supported # ── Test config ───────────────────────────────────────────────────────────── @@ -660,6 +661,55 @@ def run(idx, toks, w): expected = (("dp", "ep"),) if self.dp > 1 else ("ep",) self.assertEqual(tuple(compiled.output_shardings.spec), expected) + @unittest.skipUnless( + is_collective_stream_supported(), + "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)", + ) + def test_z_dispatch_combine_on_collective_stream(self): + """Every EP FFI custom call must carry the collective-stream annotation + so XLA schedules them on the collective stream instead of overlapping + them with other collectives.""" + T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() + dp_spec = PartitionSpec(("dp", "ep"), None) + ep_spec_3d = PartitionSpec(("dp", "ep"), None, None) + ep_spec_2d = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + + @jax.jit + def run(idx, toks, w): + idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) + toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) + w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) + recv_t, recv_w, hm, tc = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) + recv_t = jax.lax.with_sharding_constraint( + recv_t, NamedSharding(self.mesh, ep_spec_3d) + ) + recv_w = jax.lax.with_sharding_constraint( + recv_w, NamedSharding(self.mesh, ep_spec_2d) + ) + weighted = self._preweight_expert_out(recv_t, recv_w) + out = ep_combine(self.hk, hm, tc, weighted, T_dp, out_sharding=(("dp", "ep"), None)) + return jax.lax.with_sharding_constraint(out, NamedSharding(self.mesh, dp_spec)) + + hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text() + + # Every te_ep_* FFI custom call must carry the collective-stream + # annotation so XLA places it on the collective stream. + ep_lines = [l for l in hlo.splitlines() if 'custom_call_target="te_ep_' in l] + self.assertTrue(ep_lines, f"no te_ep_* custom calls in compiled HLO:\n{hlo}") + missing = [ + l.strip()[:200] + for l in ep_lines + if '_xla_stream_annotation="collective"' not in l.replace(" ", "") + ] + self.assertFalse( + missing, + "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing), + ) + def test_z_no_unexpected_reshard_in_hlo_bwd(self): """Compiled bwd HLO must not insert XLA collectives outside the EP FFI.""" T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 77e60afbcd..806e7ae480 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -24,6 +24,18 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive from ..sharding import global_mesh_resource, get_mesh_axis_size +from ..version_utils import is_collective_stream_supported + + +def _on_collective_stream(func): + """Pin ``func``'s ops to XLA's collective stream so the scheduler serializes + them with native collectives. No-op on JAX that lacks the annotation.""" + if not is_collective_stream_supported(): + return func + from jax.experimental.compute_on import compute_on + + return compute_on("gpu_stream:collective")(func) # pylint: disable=not-callable + __all__ = [ "EpConfig", @@ -894,6 +906,7 @@ def shardy_sharding_rule(*args): # ── Public-ish helpers (used by jax/ep.py) ────────────────────────────────── +@_on_collective_stream def ep_prepare(cfg: EpLayerConfig, topk_idx): """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``.""" return EpPreparePrimitive.outer_primitive.bind( @@ -904,6 +917,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx): ) +@_on_collective_stream def ep_dispatch_fwd( cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ): @@ -920,6 +934,7 @@ def ep_dispatch_fwd( ) +@_on_collective_stream def ep_combine_fwd( cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None ): @@ -935,6 +950,7 @@ def ep_combine_fwd( ) +@_on_collective_stream def ep_dispatch_bwd( cfg: EpLayerConfig, handle_mem, @@ -956,6 +972,7 @@ def ep_dispatch_bwd( ) +@_on_collective_stream def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank): """Backward of combine; returns grad_expert_out [num_procs, recv_capacity_per_rank, H].""" return EpCombineBwdPrimitive.outer_primitive.bind( diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e4619d8670..8765a4fce0 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -64,6 +64,16 @@ def is_triton_autotuned_alias_safe() -> bool: return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR) +# XLA gained the ``gpu_stream:collective`` stream annotation in openxla/xla#39604, +# which ships in the JAX 0.10.0 release. Older XLA fatally fails on it. +_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.0" + + +def is_collective_stream_supported() -> bool: + """Return True if the installed JAX supports the gpu_stream:collective annotation.""" + return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION) + + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -77,6 +87,7 @@ def is_triton_extension_supported() -> bool: __all__ = [ "jax_version_meet_requirement", "is_triton_autotuned_alias_safe", + "is_collective_stream_supported", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION",