From 8e60015d954c83a02ff68a74dad4ed175516d9b9 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 10 Nov 2026 13:56:10 -0800 Subject: [PATCH] [JAX] Save routing state when checkpointing EP dispatch Save the post-dispatch routing handle and expert token counts under the dispatch checkpoint name so backward can reuse dispatch state without recomputing communication. Add regression coverage for saved versus recomputed dispatch, gradients, scan handle relocation, empty experts, and overflow backward. Validate split and fused prepare/dispatch paths. Signed-off-by: Jeremy Berchtold --- tests/jax/test_multi_process_ep.py | 52 +++++++++ tests/jax/test_te_ep_moe.py | 162 +++++++++++++++++++++++++++++ transformer_engine/jax/moe.py | 6 +- 3 files changed, 219 insertions(+), 1 deletion(-) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 38822812372..de0ad6ad6c3 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -1142,6 +1142,58 @@ def run(idx, toks, w): # The rank owning expert 0 demands more than it can receive. self.assertGreater(int(trt_prep.max()), self.recv_capacity_per_rank) + def test_checkpointed_routing_state_with_overflow(self): + """A saved routing handle remains usable by combine backward after drops.""" + from jax.ad_checkpoint import checkpoint_name + + total_tokens, indices, tokens, weights = self._make_concentrated_inputs() + token_spec = PartitionSpec(("dp", "ep"), None) + recv_spec = PartitionSpec(("dp", "ep"), None, None) + + def loss_fn(scale, idx, toks, routing_weights): + idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, token_spec)) + toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, token_spec)) + routing_weights = jax.lax.with_sharding_constraint( + routing_weights, NamedSharding(self.mesh, token_spec) + ) + received, received_weights, handle, counts, total_recv = ep_dispatch( + self.hk, idx, toks, routing_weights, self.recv_capacity_per_rank + ) + received, received_weights, handle, counts = ( + checkpoint_name(value, "saved_dispatch") + for value in (received, received_weights, handle, counts) + ) + align = max(self.hk.dispatch_output_per_expert_alignment, 1) + padded = ((counts + align - 1) // align) * align + valid_slots = jnp.sum(padded, axis=-1, keepdims=True) + mask = jnp.arange(self.recv_capacity_per_rank)[None, :] < valid_slots + received = jnp.where(mask[..., None], received, 0) + received_weights = jnp.where(mask, received_weights, 0) + expert_output = (received * received_weights[..., None] * scale).astype(toks.dtype) + expert_output = jax.lax.with_sharding_constraint( + expert_output, NamedSharding(self.mesh, recv_spec) + ) + output = ep_combine(self.hk, handle, counts, expert_output, total_tokens) + return jnp.sum(output.astype(jnp.float32) ** 2), total_recv + + checkpointed = jax.checkpoint( + loss_fn, policy=jax.checkpoint_policies.save_only_these_names("saved_dispatch") + ) + with self.mesh, global_shard_guard(self.mr): + args = (jnp.asarray(0.5, jnp.float32), indices, tokens, weights) + expected = jax.jit(jax.value_and_grad(loss_fn, has_aux=True))(*args) + actual = jax.jit(jax.value_and_grad(checkpointed, has_aux=True))(*args) + jax.block_until_ready((actual, expected)) + actual = jmu.process_allgather(actual, tiled=True) + expected = jmu.process_allgather(expected, tiled=True) + for value, reference in zip( + jax.tree_util.tree_leaves(actual), jax.tree_util.tree_leaves(expected) + ): + self.assertTrue(np.all(np.isfinite(value))) + np.testing.assert_allclose(value, reference, atol=1e-5, rtol=1e-5) + self.assertGreater(int(np.max(np.asarray(actual[0][1]))), self.recv_capacity_per_rank) + self.assertGreater(float(np.max(np.abs(np.asarray(actual[1])))), 0.0) + # lax.scan handle_mem relocation diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index ad5c897b8d5..89abfe6c265 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -45,6 +45,7 @@ import importlib import os +import re os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") @@ -1139,6 +1140,167 @@ def recorded_checkpoint_name(value, name): assert np.all(np.isfinite(_to_global_numpy(_unwrap(grads["params"][name]), mesh))) +def _assert_checkpoint_parity(actual, expected, mesh, *, relative_tolerance=1e-3): + """Compare saved-state and recomputed executions of the same quantized block.""" + actual_leaves, actual_tree = jax.tree_util.tree_flatten(actual) + expected_leaves, expected_tree = jax.tree_util.tree_flatten(expected) + assert actual_tree == expected_tree + for actual_leaf, expected_leaf in zip(actual_leaves, expected_leaves): + actual_np = _to_global_numpy(actual_leaf, mesh) + expected_np = _to_global_numpy(expected_leaf, mesh) + assert np.all(np.isfinite(actual_np)) + if jnp.issubdtype(actual_leaf.dtype, jnp.integer): + np.testing.assert_array_equal(actual_np, expected_np) + else: + # Fused EP permits different within-expert row orders. In particular, + # weight-gradient reductions need not be bitwise reproducible. + np.testing.assert_allclose( + actual_np.astype(np.float32), + expected_np.astype(np.float32), + atol=1e-7, + rtol=1e-3, + ) + # Small gradients must not pass solely because of the absolute + # tolerance: losing or misrouting an entire gradient is a failure. + actual_f64 = actual_np.astype(np.float64) + expected_f64 = expected_np.astype(np.float64) + reference_norm = np.linalg.norm(expected_f64.reshape(-1)) + if reference_norm > 0: + error_norm = np.linalg.norm((actual_f64 - expected_f64).reshape(-1)) + assert error_norm <= relative_tolerance * reference_norm, ( + f"Checkpoint relative L2 error {error_norm / reference_norm:.3%} " + f"for shape {actual_np.shape}" + ) + + +def _forward_dispatch_count(executable): + """Count forward EP custom calls, excluding the required backward transport.""" + return len( + re.findall( + r'custom_call_target="te_ep_(?:dispatch|prepare_and_dispatch)_ffi"', + executable.as_text(), + ) + ) + + +class TestTeEpMoeDispatchCheckpoint: + """Checkpoint the complete dispatch state, including opaque routing metadata.""" + + @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) + @pytest.mark.parametrize("skewed", [False, True], ids=["balanced", "empty_experts"]) + def test_saved_dispatch_removes_redispatch(self, mesh, make_block, quantization, skewed): + recipe = _quantization_recipe(quantization) + block = make_block( + quantization_recipe=recipe, + dispatch_checkpoint_name="saved_dispatch", + score_function="sigmoid" if skewed else "softmax", + use_expert_routing_bias=skewed, + expert_bias_init=_strong_expert_bias_init if skewed else None, + ) + x = _make_inputs(jax.random.PRNGKey(40)) + variables, _, _, counts = _init_apply(block, mesh, x, jax.random.PRNGKey(41)) + if skewed: + assert np.any(_to_global_numpy(counts, mesh) == 0) + + def loss_fn(params, inputs): + output, _, total_recv, expert_counts = block.apply(params, inputs) + return jnp.mean(output.astype(jnp.float32) ** 2), (output, total_recv, expert_counts) + + saved = jax.checkpoint( + loss_fn, policy=jax.checkpoint_policies.save_only_these_names("saved_dispatch") + ) + recomputed = jax.checkpoint(loss_fn, policy=jax.checkpoint_policies.nothing_saveable) + with _ctx(mesh): + x = _shard_inputs(x, mesh) + reference = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True)) + saved_exe = ( + jax.jit(jax.value_and_grad(saved, argnums=(0, 1), has_aux=True)) + .lower(variables, x) + .compile() + ) + recomputed_exe = ( + jax.jit(jax.value_and_grad(recomputed, argnums=(0, 1), has_aux=True)) + .lower(variables, x) + .compile() + ) + expected = reference(variables, x) + actual = saved_exe(variables, x) + control = recomputed_exe(variables, x) + jax.block_until_ready((expected, actual, control)) + # One forward dispatch per layer. The control proves that the assertion + # distinguishes recomputation from merely attaching checkpoint names. + assert _forward_dispatch_count(saved_exe) == 1 + assert _forward_dispatch_count(recomputed_exe) == 2 + prepare_pattern = r'custom_call_target="te_ep_prepare_ffi"' + expected_prepares = 0 if os.environ.get("NVTE_EP_FUSED_PREPARE_DISPATCH") == "1" else 1 + assert len(re.findall(prepare_pattern, saved_exe.as_text())) == expected_prepares + assert len(re.findall(prepare_pattern, recomputed_exe.as_text())) == 2 * expected_prepares + _assert_checkpoint_parity(actual, expected, mesh) + _assert_checkpoint_parity(control, expected, mesh) + + @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) + def test_saved_handles_in_scan(self, mesh, make_block, quantization): + recipe = _quantization_recipe(quantization) + block = make_block( + quantization_recipe=recipe, + dispatch_checkpoint_name="saved_dispatch", + ) + x = _make_inputs(jax.random.PRNGKey(42)) + variables, _, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(43)) + + def layer(params, inputs): + output, _, _, _ = block.apply(params, inputs) + return output + + saved_layer = jax.checkpoint( + layer, policy=jax.checkpoint_policies.save_only_these_names("saved_dispatch") + ) + + def scanned_loss(params, inputs, *, apply_layer): + def body(carry, layer_input): + return carry, apply_layer(params, layer_input) + + _, outputs = jax.lax.scan(body, None, inputs) + return jnp.mean(outputs.astype(jnp.float32) ** 2), outputs + + def unrolled_loss(params, inputs): + outputs = jnp.stack([layer(params, inputs[i]) for i in range(2)]) + return jnp.mean(outputs.astype(jnp.float32) ** 2), outputs + + scanned = jax.jit( + jax.value_and_grad( + partial(scanned_loss, apply_layer=saved_layer), argnums=(0, 1), has_aux=True + ) + ) + plain_scanned = jax.jit( + jax.value_and_grad( + partial(scanned_loss, apply_layer=layer), argnums=(0, 1), has_aux=True + ) + ) + unrolled = jax.jit(jax.value_and_grad(unrolled_loss, argnums=(0, 1), has_aux=True)) + # Different inputs give each layer its own routing state. Reuse the same + # executables on another batch to expose stale handles across invocations. + for seed in (44, 45): + with _ctx(mesh): + inputs = jnp.stack( + [ + _shard_inputs(_make_inputs(jax.random.PRNGKey(seed + i)), mesh) + for i in range(2) + ] + ) + actual = scanned(variables, inputs) + control = plain_scanned(variables, inputs) + expected = unrolled(variables, inputs) + jax.block_until_ready((actual, control, expected)) + _assert_checkpoint_parity(actual, control, mesh) + # Scan and unrolled forms accumulate shared BF16 parameter gradients + # differently. Use the same scan as the checkpointing control, then + # allow one BF16 epsilon for the independent unrolled cross-check. + _assert_checkpoint_parity( + actual, expected, mesh, relative_tolerance=float(jnp.finfo(DTYPE).eps) + ) + + def test_explicitly_disabled_fusion_with_native_weights( mesh, monkeypatch, make_block, use_cudnn_fusion ): diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 4e9bfe4aa8d..15379b8846d 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -1561,6 +1561,9 @@ def _moe_fwd_rule( if dispatch_checkpoint_name is not None: recv_tokens = checkpoint_name(recv_tokens, dispatch_checkpoint_name) recv_topk_weights = checkpoint_name(recv_topk_weights, dispatch_checkpoint_name) + # Backward needs routing state as well as the dispatched activations. + handle_mem = checkpoint_name(handle_mem, dispatch_checkpoint_name) + token_counts = checkpoint_name(token_counts, dispatch_checkpoint_name) recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3_spec)) recv_topk_weights = jax.lax.with_sharding_constraint( recv_topk_weights, NamedSharding(mesh, ep2_spec) @@ -2160,7 +2163,8 @@ def moe( ``None`` leaves the value unnamed. dispatch_checkpoint_name : Optional[str] JAX rematerialization checkpoint name for the EP dispatch outputs. - ``None`` leaves these values unnamed; the opaque prepare handle is never named. + Includes the opaque routing handle and expert token counts needed by backward. + ``None`` leaves these values unnamed. mesh_resource : Optional[MeshResource] Physical parallelism resources. ``None`` uses the active ``global_shard_guard`` context; an explicit resource takes precedence.