From 2abab4d8b54d5f4ea132ab77ab61e791828bea0b Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 24 Jul 2026 16:13:20 -0700 Subject: [PATCH] [JAX] Add NCCL EP drop-on-overflow policy and total_recv_tokens output Signed-off-by: Phuong Nguyen --- examples/jax/ep/bench/ep_bench.py | 4 +- examples/jax/ep/ep_moe.py | 2 +- tests/jax/test_multi_process_ep.py | 165 ++++++++++++++++-- tests/jax/test_te_ep_moe.py | 6 +- transformer_engine/jax/cpp_extensions/ep.py | 27 ++- transformer_engine/jax/csrc/extensions.h | 3 +- transformer_engine/jax/csrc/extensions/ep.cpp | 19 +- .../jax/csrc/extensions/pybind.cpp | 4 +- transformer_engine/jax/ep.py | 34 +++- transformer_engine/jax/flax/moe.py | 5 +- transformer_engine/jax/moe.py | 20 ++- 11 files changed, 237 insertions(+), 52 deletions(-) diff --git a/examples/jax/ep/bench/ep_bench.py b/examples/jax/ep/bench/ep_bench.py index 27ad8ca146..50b95e34f3 100644 --- a/examples/jax/ep/bench/ep_bench.py +++ b/examples/jax/ep/bench/ep_bench.py @@ -148,7 +148,7 @@ def main(): @jax.jit def run_prepare(idx): - tc, hm = tex_ep.ep_prepare(cfg, idx) + tc, _trt, hm = tex_ep.ep_prepare(cfg, idx) return tc, hm @jax.jit @@ -160,7 +160,7 @@ def run_dispatch(hm, idx, toks, w): @jax.jit def run_dispatch_vjp(idx, toks, w): - recv_t, recv_w, _hm, _tc = ep_dispatch(cfg, idx, toks, w, recv_capacity_per_rank) + recv_t, recv_w, _hm, _tc, _trt = ep_dispatch(cfg, idx, toks, w, recv_capacity_per_rank) recv_t = jax.lax.with_sharding_constraint(recv_t, NamedSharding(mesh, ep_spec_3d)) recv_w = jax.lax.with_sharding_constraint(recv_w, NamedSharding(mesh, ep_spec_2d)) return recv_t, recv_w diff --git a/examples/jax/ep/ep_moe.py b/examples/jax/ep/ep_moe.py index 671150d655..adf0f2ca0a 100644 --- a/examples/jax/ep/ep_moe.py +++ b/examples/jax/ep/ep_moe.py @@ -231,7 +231,7 @@ def _moe_layer(args, cfg, mesh, topk_idx, tokens, topk_w, local_kernels): local_kernels = jax.lax.with_sharding_constraint( local_kernels, NamedSharding(mesh, kernel_spec) ) - recv_tokens, recv_topk_w, handle_mem, _tc = ep_dispatch( + recv_tokens, recv_topk_w, handle_mem, _tc, _trt = ep_dispatch( cfg, topk_idx, tokens, topk_w, args.recv_capacity_per_rank ) recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3)) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index c89f1234c6..4121e5d1b4 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -14,6 +14,9 @@ - ``ep_dispatch`` custom_vjp: exact per-(t, k) ``grad_topk_weights`` under skewed upstream gradients (no k-axis averaging). - HLO reshard guard: compile-only, no XLA collectives outside the EP FFI. + - Drop-on-overflow: ``ep_finalize`` + re-bootstrap with a small recv capacity + drops the overflow instead of trapping; ``total_recv_tokens`` reports the + pre-drop demand. Launch via tests/jax/multi_process_launch_ep.sh (one process per GPU). """ @@ -33,6 +36,7 @@ from transformer_engine.jax.ep import ( EpLayerConfig, ep_bootstrap, + ep_finalize, ep_dispatch, ep_combine, _ep_domain_for_rank, @@ -238,8 +242,8 @@ def test_two_handle_mems_no_aliasing(self): @jax.jit def run(idx): - _tc_a, ha = ep_prepare(ka, idx) - _tc_b, hb = ep_prepare(kb, idx) + _tc_a, _trt_a, ha = ep_prepare(ka, idx) + _tc_b, _trt_b, hb = ep_prepare(kb, idx) return ha, hb hm_a, hm_b = run(idx_s) @@ -266,7 +270,9 @@ def test_two_layer_dispatch_no_handle_aliasing(self): w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) def one_layer(hk, idx, toks, w_): - recv_t, recv_w, hm, tc = ep_dispatch(hk, idx, toks, w_, self.recv_capacity_per_rank) + recv_t, recv_w, hm, tc, _trt = ep_dispatch( + hk, idx, toks, w_, self.recv_capacity_per_rank + ) recv_t = jax.lax.with_sharding_constraint( recv_t, NamedSharding(self.mesh, ep_spec_3d) ) @@ -304,23 +310,33 @@ def run(idx, ta_, tb_, w_): ) def test_primitive_prepare(self): - """ep_prepare returns token_counts and handle_mem of the expected shapes.""" + """ep_prepare returns token_counts, total_recv_tokens and handle_mem. + + total_recv_tokens is the per-rank pre-drop recv-slot total; with no + overflow it equals the padded per-expert count sum that dispatch fills. + """ T_global, topk_idx, _tokens, _w = self._make_identity_inputs() del T_global dp_spec = PartitionSpec(("dp", "ep"), None) + align = max(int(self.hk.dispatch_output_per_expert_alignment), 1) with self.mesh, global_shard_guard(self.mr): idx_s = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) @jax.jit def run(idx): - tc, hm = ep_prepare(self.hk, idx) - return tc, hm + tc, trt, hm = ep_prepare(self.hk, idx) + return tc, trt, hm - tc, hm = run(idx_s) - tc.block_until_ready() + tc, trt, hm = run(idx_s) + trt.block_until_ready() self.assertEqual(tc.shape, (self.dp * self.ep, NUM_LOCAL_EXPERTS)) self.assertEqual(hm.shape[0], self.dp * self.ep) self.assertGreater(hm.shape[1], 0) + self.assertEqual(trt.shape, (self.dp * self.ep, 1)) + self.assertEqual(trt.dtype, jnp.int32) + padded = ((np.asarray(tc).astype(np.int64) + align - 1) // align) * align + expected = padded.sum(axis=-1, keepdims=True) + np.testing.assert_array_equal(np.asarray(trt).astype(np.int64), expected) def _run_identity_round_trip(self, nonuniform): T_global, topk_idx, tokens, topk_w = self._make_identity_inputs(nonuniform=nonuniform) @@ -335,7 +351,7 @@ def _run_identity_round_trip(self, nonuniform): @jax.jit def run(idx, toks, w): - _tc, hm = ep_prepare(self.hk, idx) + _tc, _trt, hm = ep_prepare(self.hk, idx) recv_t, recv_w = ep_dispatch_fwd( self.hk, hm, idx, toks, w, self.recv_capacity_per_rank ) @@ -401,7 +417,7 @@ def loss_fn(toks): toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) - recv_t, recv_w, hm, tc = ep_dispatch( + recv_t, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_t = jax.lax.with_sharding_constraint( @@ -452,7 +468,7 @@ def test_dispatch_combine_3d_input_output(self): @jax.jit def run(idx, toks, w): - recv_t, recv_w, hm, _tc = ep_dispatch( + recv_t, recv_w, hm, _tc, _trt = 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_t)) @@ -503,7 +519,7 @@ def loss_fn(toks): toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) - recv_tokens, _recv_w, _hm, tc = ep_dispatch( + recv_tokens, _recv_w, _hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_tokens = jax.lax.with_sharding_constraint( @@ -557,7 +573,7 @@ def loss_fn(eo): toks = jax.lax.with_sharding_constraint(tokens, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) - _recv_tokens, recv_w, hm, tc = ep_dispatch( + _recv_tokens, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_w = jax.lax.with_sharding_constraint( @@ -603,7 +619,7 @@ def loss_fn(idx_in, tok_in, w_in): idx_in = jax.lax.with_sharding_constraint(idx_in, NamedSharding(self.mesh, dp_spec)) tok_in = jax.lax.with_sharding_constraint(tok_in, NamedSharding(self.mesh, dp_spec)) w_in = jax.lax.with_sharding_constraint(w_in, NamedSharding(self.mesh, dp_spec)) - _recv_t, recv_w, _h, _tc = ep_dispatch( + _recv_t, recv_w, _h, _tc, _trt = ep_dispatch( self.hk, idx_in, tok_in, w_in, self.recv_capacity_per_rank ) # Per-slot index scale ⇒ each slot's contribution differs. @@ -644,7 +660,7 @@ 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( + recv_t, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_t = jax.lax.with_sharding_constraint( @@ -688,7 +704,7 @@ 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( + recv_t, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_t = jax.lax.with_sharding_constraint( @@ -739,7 +755,9 @@ def fwd(eo, toks, idx, w): toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) - _rt, rw, hm, tc = ep_dispatch(self.hk, idx, toks, w, self.recv_capacity_per_rank) + _rt, rw, hm, tc, _trt = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) rw = jax.lax.with_sharding_constraint(rw, NamedSharding(self.mesh, ep_spec_2d)) weighted = self._preweight_expert_out(eo, rw) combined = ep_combine( @@ -769,6 +787,117 @@ def bwd_only(eo, toks, idx, w, g): self.assertEqual(hlo.count(op), 0, f"unexpected XLA {op} in bwd HLO:\n{hlo}") +# ── Drop-on-overflow ───────────────────────────────────────────────────────── + + +class TestEPOverflowDrop(unittest.TestCase): + """Re-bootstraps with drop_on_overflow=True and a small recv capacity. + + ``ep_finalize`` tears down the default TestEP communicator so this class can + re-bootstrap with its own config in the same process. Every token routes its + top-1 slot to expert 0, so the rank owning it demands more than the recv + capacity; the dispatch drops the excess instead of trapping, while + total_recv_tokens still reports the pre-drop demand. + """ + + ALIGN = 16 + # Each EP group routes all OVF_TOKENS_PER_DP_SHARD top-1 slots to expert 0, + # whose padded count then exceeds the recv capacity. HT mode requires the + # recv capacity to be >= the per-rank dispatch count (tokens // ep), so the + # capacity sits between that floor and the concentrated expert-0 demand. + OVF_TOKENS_PER_DP_SHARD = 48 + OVF_RECV_CAPACITY = NUM_LOCAL_EXPERTS * ALIGN # 32 slots per rank + + @classmethod + def setUpClass(cls): + sm = _local_device_sm() + if sm is not None and sm < 90: + raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{sm})") + cls.num_procs = jax.process_count() + cls.rank = jax.process_index() + cls.dp, cls.ep = _factor_dp_ep(cls.num_procs) + cls.num_experts = NUM_LOCAL_EXPERTS * cls.ep + cls.recv_capacity_per_rank = cls.OVF_RECV_CAPACITY + # True per-rank dispatch count; HT mode needs recv_capacity >= this. + cls.max_tokens_per_rank = cls.OVF_TOKENS_PER_DP_SHARD // cls.ep + cls.mesh = _build_mesh(cls.dp, cls.ep) + cls.mr = MeshResource(dp_resource="dp", ep_resource="ep") + # Drop any communicator a prior test class bootstrapped, then re-init. + ep_finalize() + with cls.mesh, global_shard_guard(cls.mr): + ep_bootstrap( + world_size=cls.num_procs, + rank=cls.rank, + num_experts=cls.num_experts, + max_tokens_per_rank=cls.max_tokens_per_rank, + recv_capacity_per_rank=cls.recv_capacity_per_rank, + hidden_dim=HIDDEN_DIM, + drop_on_overflow=True, + ) + cls.hk = EpLayerConfig(top_k=TOP_K, dispatch_output_per_expert_alignment=cls.ALIGN) + + @classmethod + def tearDownClass(cls): + # Leave a clean slate so another class can bootstrap after us. + ep_finalize() + + def _make_concentrated_inputs(self): + """All top-1 routes to expert 0; top-2 spread over the rest, so the rank + owning expert 0 overloads beyond recv_capacity_per_rank.""" + T_global = self.OVF_TOKENS_PER_DP_SHARD * self.dp + E = self.num_experts + topk_idx = np.empty((T_global, TOP_K), dtype=np.int32) + for t in range(T_global): + topk_idx[t, 0] = 0 + for k in range(1, TOP_K): + topk_idx[t, k] = 1 + (t % (E - 1)) + topk_idx = jnp.asarray(topk_idx) + topk_w = jnp.full((T_global, TOP_K), 1.0 / TOP_K, dtype=jnp.float32) + tokens = jnp.asarray( + np.linspace(0.1, 0.9, T_global * HIDDEN_DIM, dtype=np.float32).reshape( + T_global, HIDDEN_DIM + ), + dtype=jnp.bfloat16, + ) + return T_global, topk_idx, tokens, topk_w + + def test_overflow_drops_without_trap(self): + """Dispatch drops the overflow instead of trapping; total_recv_tokens + reports the pre-drop demand, which exceeds recv_capacity_per_rank.""" + _T, topk_idx, tokens, topk_w = self._make_concentrated_inputs() + dp_spec = PartitionSpec(("dp", "ep"), None) + ep_spec_3d = PartitionSpec(("dp", "ep"), None, 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)) + _tc, trt_prep, _hm = ep_prepare(self.hk, idx) + recv_t, _rw, _hm2, _tc2, trt_disp = 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) + ) + return trt_prep, trt_disp, recv_t + + trt_prep, trt_disp, recv_t = run(topk_idx, tokens, topk_w) + recv_t.block_until_ready() + + # Reaching here means the dispatch dropped the overflow rather than + # trapping (an overflow trap is a device-side abort). + self.assertEqual(recv_t.shape[1], self.recv_capacity_per_rank) + # total_recv_tokens is sharded across processes; gather before host reads. + trt_prep = np.asarray(jmu.process_allgather(trt_prep, tiled=True)).reshape(-1) + trt_disp = np.asarray(jmu.process_allgather(trt_disp, tiled=True)).reshape(-1) + # prepare and dispatch see the same routing -> identical pre-drop totals. + np.testing.assert_array_equal(trt_prep, trt_disp) + # The rank owning expert 0 demands more than it can receive. + self.assertGreater(int(trt_prep.max()), self.recv_capacity_per_rank) + + # ── EP domain grouping (single-process; runs under plain pytest) ───────────── @@ -819,7 +948,7 @@ def test_ep_tp_splits_domains(self): ) loader = unittest.TestLoader() - test_cases = (TestEP, TestEpDomainGrouping) + test_cases = (TestEP, TestEPOverflowDrop, TestEpDomainGrouping) target = os.environ.get("TARGET_TEST") if target: name = target.split(".")[-1] diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index d08765e184..55ed415d21 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -429,7 +429,7 @@ def _init_apply(block, mesh, x, key): x_sh = _shard_inputs(x, mesh) variables = jax.jit(block.init)(key, x_sh) jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) - output, aux = jax.jit(block.apply)(variables, x_sh) + output, aux, _trt = jax.jit(block.apply)(variables, x_sh) jax.block_until_ready(output) return variables, output, aux @@ -445,7 +445,7 @@ def _grad_step(block, variables, mesh, x, *, include_aux=False): x_sh = _shard_inputs(x, mesh) def loss_fn(variables, x): - output, aux = block.apply(variables, x) + output, aux, _trt = block.apply(variables, x) loss = jnp.mean(output.astype(jnp.float32) ** 2) if include_aux and aux is not None: loss = loss + aux.astype(jnp.float32) @@ -464,7 +464,7 @@ def _grad_aux_only(block, variables, mesh, x): x_sh = _shard_inputs(x, mesh) def aux_only(variables, x): - _, aux = block.apply(variables, x) + _, aux, _trt = block.apply(variables, x) return aux.astype(jnp.float32) grads = jax.jit(jax.grad(aux_only))(variables, x_sh) diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 806e7ae480..d0d472c6f2 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -42,6 +42,7 @@ def _on_collective_stream(func): "EpLayerConfig", "set_ep_config", "get_ep_config", + "reset_ep_config", "ep_handle_mem_size", "ep_prepare", "ep_dispatch_fwd", @@ -89,6 +90,12 @@ def get_ep_config() -> EpConfig: return _ep_config +def reset_ep_config() -> None: + """Clear the cached EpConfig so a later ep_bootstrap starts fresh (see ep_finalize).""" + global _ep_config + _ep_config = None + + @dataclass(frozen=True) class EpLayerConfig: """Per-layer EP config; mirrors C ``NVTEEpLayerConfig``. @@ -215,8 +222,10 @@ def abstract(topk_idx_aval, *, top_k, dispatch_output_per_expert_alignment, is_o ) leading = _ep_leading_dims(is_outer) token_counts_aval = jax.core.ShapedArray(leading + (num_local_experts,), jnp.int32) + # Per-rank pre-drop recv-slot total (includes tokens dropped on overflow). + total_recv_tokens_aval = jax.core.ShapedArray(leading + (1,), jnp.int32) handle_mem_aval = jax.core.ShapedArray(leading + (handle_mem_size,), jnp.uint8) - return token_counts_aval, handle_mem_aval + return token_counts_aval, total_recv_tokens_aval, handle_mem_aval @staticmethod def outer_abstract(*args, **kwargs): @@ -236,13 +245,13 @@ def lowering(ctx, topk_idx, *, top_k, dispatch_output_per_expert_alignment, is_o @staticmethod def impl(topk_idx, top_k, dispatch_output_per_expert_alignment, is_outer): assert EpPreparePrimitive.inner_primitive is not None - token_counts, handle_mem = EpPreparePrimitive.inner_primitive.bind( + token_counts, total_recv_tokens, handle_mem = EpPreparePrimitive.inner_primitive.bind( topk_idx, top_k=top_k, dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, is_outer=is_outer, ) - return token_counts, handle_mem + return token_counts, total_recv_tokens, handle_mem @staticmethod def batcher(batched_args, batch_dims, *, top_k, dispatch_output_per_expert_alignment, is_outer): @@ -262,9 +271,11 @@ def partition( f" with the topk dim replicated; got spec={idx_spec}." ) arg_shardings = tuple(a.sharding for a in arg_infos) - # token_counts / handle_mem inherit the input's leading axis (trailing dims auto-pad to None). + # token_counts / total_recv_tokens / handle_mem inherit the input's leading + # axis (trailing dims auto-pad to None). leading_spec = PartitionSpec(idx_spec[0]) tc_sharding = NamedSharding(mesh, leading_spec) + trt_sharding = NamedSharding(mesh, leading_spec) hm_sharding = NamedSharding(mesh, leading_spec) def sharded_impl(topk_idx): @@ -272,7 +283,7 @@ def sharded_impl(topk_idx): topk_idx, top_k, dispatch_output_per_expert_alignment, False ) - return mesh, sharded_impl, (tc_sharding, hm_sharding), arg_shardings + return mesh, sharded_impl, (tc_sharding, trt_sharding, hm_sharding), arg_shardings @staticmethod def shardy_sharding_rule(*args): @@ -281,7 +292,7 @@ def shardy_sharding_rule(*args): value_types = args[-2] topk_idx_rank = len(value_types[0].shape) in_axes = " ".join(f"L{i}" for i in range(topk_idx_rank - 1)) + " topk" - return f"{in_axes} -> EPL nle, EPL hm" + return f"{in_axes} -> EPL nle, EPL trt, EPL hm" register_primitive(EpPreparePrimitive) @@ -908,7 +919,9 @@ def shardy_sharding_rule(*args): @_on_collective_stream def ep_prepare(cfg: EpLayerConfig, topk_idx): - """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``.""" + """Exchange routing metadata for ``cfg``; return + ``(token_counts, total_recv_tokens, handle_mem)``. ``total_recv_tokens`` is + the per-rank pre-drop recv-slot total (includes tokens dropped on overflow).""" return EpPreparePrimitive.outer_primitive.bind( topk_idx, top_k=int(cfg.top_k), diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index b9c7c849f2..9bd0940c4b 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -209,7 +209,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler); // the group will dispatch. void SetEpBootstrapParams(pybind11::bytes unique_id_bytes, int ep_size, int rank_within_group, int num_experts, int max_tokens_per_rank, int max_recv_tokens_per_rank, - int hidden_dim, int max_num_sms, int max_token_dtype); + int hidden_dim, int max_num_sms, int max_token_dtype, + bool drop_on_overflow); void ReleaseEpResources(); // Return the handle_mem byte size for a layer config. size_t EpHandleMemSize(int top_k, size_t dispatch_output_per_expert_alignment); diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index cdd7730204..aa5ed27faa 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -35,6 +35,7 @@ struct EpBootstrapParams { int hidden_dim = 0; int max_num_sms = 0; NVTEDType max_token_dtype = kNVTEBFloat16; + bool drop_on_overflow = false; }; class EpResources { @@ -53,7 +54,8 @@ class EpResources { .hidden_dim = p.hidden_dim, .num_comm_sms = p.max_num_sms, .max_token_dtype = p.max_token_dtype, - .zero_copy = 0}; + .zero_copy = 0, + .drop_on_overflow = p.drop_on_overflow}; try { nvte_ep_initialize(static_cast(comm_), &cfg); } catch (...) { @@ -125,7 +127,8 @@ struct EpConfig { // synchronize via the UID broadcast). void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int rank_within_group, int num_experts, int max_tokens_per_rank, int max_recv_tokens_per_rank, - int hidden_dim, int max_num_sms, int max_token_dtype) { + int hidden_dim, int max_num_sms, int max_token_dtype, + bool drop_on_overflow) { std::string uid_str = unique_id_bytes_obj; NVTE_CHECK(static_cast(uid_str.size()) >= 128, "unique_id_bytes must be at least 128 bytes (ncclUniqueId size)."); @@ -143,6 +146,7 @@ void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int g_ep_params.hidden_dim = hidden_dim; g_ep_params.max_num_sms = max_num_sms; g_ep_params.max_token_dtype = static_cast(max_token_dtype); + g_ep_params.drop_on_overflow = drop_on_overflow; g_ep_params_set = true; } // Acquire outside the lock: EpResources ctor runs ncclCommInitRank which is @@ -196,8 +200,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpInstantiateHandler, EpInstantiateImpl, FFI::Bind // ── ep_prepare ──────────────────────────────────────────────────────────────── Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type topk_idx, - Result_Type recv_tokens_per_expert, Result_Type handle_mem, - EpConfig config) { + Result_Type recv_tokens_per_expert, Result_Type total_recv_tokens, + Result_Type handle_mem, EpConfig config) { (void)ep_state; // lifetime only. auto topk_dims = topk_idx.dimensions(); NVTE_CHECK(topk_dims.size() >= 2, @@ -215,6 +219,10 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T auto recv_tokens_per_expert_ = TensorWrapper(recv_tokens_per_expert->untyped_data(), tc_shape, DType::kInt32); + std::vector trt_shape = {static_cast(total_recv_tokens->element_count())}; + auto total_recv_tokens_ = + TensorWrapper(total_recv_tokens->untyped_data(), trt_shape, DType::kInt32); + std::vector hm_shape = {static_cast(handle_mem->element_count())}; auto handle_mem_ = TensorWrapper(handle_mem->untyped_data(), hm_shape, DType::kByte); @@ -223,7 +231,7 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T .dispatch_output_per_expert_alignment = static_cast(config.dispatch_output_per_expert_alignment)}; nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), recv_tokens_per_expert_.data(), - /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); + total_recv_tokens_.data(), &layer_cfg, stream); return ffi_with_cuda_error_check(); } @@ -233,6 +241,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, .Ctx<::xla::ffi::State>() // EP state .Arg() // topk_idx .Ret() // recv_tokens_per_expert + .Ret() // total_recv_tokens .Ret() // handle_mem .Attrs(), FFI_CudaGraph_Traits); diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index db5468afe6..3927e2686e 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -159,8 +159,8 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("set_ep_bootstrap_params", &SetEpBootstrapParams, pybind11::arg("unique_id_bytes"), pybind11::arg("ep_size"), pybind11::arg("rank_within_group"), pybind11::arg("num_experts"), pybind11::arg("max_tokens_per_rank"), pybind11::arg("max_recv_tokens_per_rank"), - pybind11::arg("hidden_dim"), pybind11::arg("max_num_sms"), - pybind11::arg("max_token_dtype")); + pybind11::arg("hidden_dim"), pybind11::arg("max_num_sms"), pybind11::arg("max_token_dtype"), + pybind11::arg("drop_on_overflow")); m.def("release_ep_resources", &ReleaseEpResources); m.def("ep_handle_mem_size", &EpHandleMemSize, pybind11::arg("top_k"), pybind11::arg("dispatch_output_per_expert_alignment") = 0); diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index db91ea13a4..32074b49e4 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -31,6 +31,7 @@ __all__ = [ "EpLayerConfig", "ep_bootstrap", + "ep_finalize", "ep_handle_mem_size", "ep_prepare", "ep_dispatch", @@ -108,6 +109,7 @@ def ep_bootstrap( hidden_dim, max_token_dtype=jnp.bfloat16, max_num_sms=0, + drop_on_overflow=False, ): """Initialize the EP communicator. Call once per process before any EP op. @@ -126,6 +128,9 @@ def ep_bootstrap( hidden_dim: Feature dimension of token tensors passed to ep_dispatch. max_token_dtype: Widest dtype the group will dispatch (only bfloat16 supported). max_num_sms: SM budget for EP kernels; 0 = auto. + drop_on_overflow: Drop tokens exceeding recv_capacity_per_rank instead of + trapping on overflow. Dropped tokens are still counted in + total_recv_tokens, so callers can detect overflow from it. """ if jnp.dtype(max_token_dtype) != jnp.bfloat16: raise NotImplementedError( @@ -197,6 +202,7 @@ def ep_bootstrap( hidden_dim, max_num_sms=int(max_num_sms), max_token_dtype=int(jax_dtype_to_te_dtype(max_token_dtype)), + drop_on_overflow=bool(drop_on_overflow), ) # Release the C++ anchor at interpreter shutdown so RAII can tear down NCCL. @@ -220,6 +226,20 @@ def ep_bootstrap( ) +def ep_finalize(): + """Tear down the EP communicator so ``ep_bootstrap`` can run again. + + Only for killing and re-bootstrapping EP mid-program (e.g. tests sweeping + configs); a normal run bootstraps once and lets atexit clean up. Calls the + process-global ``jax.clear_caches()`` so every cached executable releases + the NCCL comm it pins, then frees the EP resources. Call outside any active + EP computation. + """ + jax.clear_caches() + transformer_engine_jax.release_ep_resources() + tex.ep.reset_ep_config() + + def _default_out_partition_spec(): """Leading-axis default: ``(("dp","ep"),)`` if DP/FSDP is set, else ``("ep",)``.""" gsr = global_mesh_resource() @@ -243,8 +263,14 @@ def ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): ``cfg`` (the pointer-keyed C++ cache keys on handle_mem, not on cfg). Inputs are ``[..., H]`` with only the leading dim sharded as ``ep`` or ``(dp, ep)``. Returns - ``(recv_tokens, recv_topk_weights, handle_mem, token_counts)``; pass - ``handle_mem`` and ``token_counts`` to the matching ``ep_combine``. + ``(recv_tokens, recv_topk_weights, handle_mem, token_counts, total_recv_tokens)``; + pass ``handle_mem`` and ``token_counts`` to the matching ``ep_combine``. + + ``total_recv_tokens`` is the per-rank pre-drop recv-slot count (a + ``[num_procs, 1]`` sharded array); it counts dropped tokens too when + ``drop_on_overflow`` is set. When ``recv_capacity_per_rank`` is not sized for + the worst case, detect overflow by ``process_allgather``-ing it, then + ``max(...) > recv_capacity_per_rank`` flags an overflowing step. """ return _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank)[0] @@ -254,12 +280,12 @@ def _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): raise TypeError( f"ep_dispatch: topk_weights must be a floating dtype; got {topk_weights.dtype}." ) - token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx) + token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( cfg, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ) out_leading = tuple(tokens.shape[:-1]) - primal = (recv_tokens, recv_topk_weights, handle_mem, token_counts) + primal = (recv_tokens, recv_topk_weights, handle_mem, token_counts, total_recv_tokens) return primal, (handle_mem, out_leading) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 3629346e33..c73bad33be 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -169,7 +169,7 @@ def __post_init__(self): super().__post_init__() @nn.compact - def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: + def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: """Run the MoE forward pass. Parameters @@ -184,6 +184,9 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: aux_loss : Optional[jnp.ndarray] Scalar load-balancing loss when ``aux_loss_coeff > 0``, else ``None``. + total_recv_tokens : jnp.ndarray + Non-differentiable per-rank pre-drop recv-slot total; flags + overflow when ``drop_on_overflow`` is set at ep_bootstrap. """ assert ( inputs.ndim == 3 diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 887e005de6..a49d4f80c2 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -637,7 +637,7 @@ def _moe_fwd_rule( top_k=K, dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) - token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr ) @@ -790,7 +790,8 @@ def _body(*args): "x_shape": x.shape, "recv_pr": recv_pr, } - return (output, aux_loss), (ctx, static) + # total_recv_tokens is a non-differentiable overflow signal (see moe()). + return (output, aux_loss, total_recv_tokens), (ctx, static) def _moe_bwd_rule( @@ -818,7 +819,8 @@ def _moe_bwd_rule( del num_groups, group_topk, dtype # captured in residuals / unused in bwd from jax.experimental.shard_map import shard_map - d_output, d_aux_loss = cotangents + # total_recv_tokens is a non-differentiable output; its cotangent is unused. + d_output, d_aux_loss, _d_total_recv_tokens = cotangents ctx, static = residuals has_bias = static["has_bias"] @@ -1155,11 +1157,13 @@ def moe( wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), dtype: jnp.dtype = jnp.float32, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: """Run a full MoE block under a single fused custom_vjp on the TE EP path. - Returns ``(output, aux_loss)``. ``aux_loss`` is ``None`` when - ``aux_loss_coeff == 0`` and a 0-d scalar otherwise. + Returns ``(output, aux_loss, total_recv_tokens)``. ``aux_loss`` is ``None`` + when ``aux_loss_coeff == 0``, else a 0-d scalar. ``total_recv_tokens`` is a + non-differentiable pre-drop recv-slot total (grad ``None``); see + ``ep_dispatch`` for using it to detect overflow. Parameters ---------- @@ -1238,7 +1242,7 @@ def moe( else: expert_bias_arg = expert_bias.astype(jnp.float32) - output, aux_loss = _moe( + output, aux_loss, total_recv_tokens = _moe( x, gate_kernel, wi_0, @@ -1269,4 +1273,4 @@ def moe( if aux_loss_coeff <= 0.0: aux_loss = None assert output.dtype == x.dtype, f"moe() output dtype {output.dtype} != input dtype {x.dtype}" - return output, aux_loss + return output, aux_loss, total_recv_tokens