From 222ef6fec46be0207b42beeb9b63859b7ac61f49 Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Tue, 4 Aug 2026 18:11:39 +0000 Subject: [PATCH 1/4] feat: implement token activation deduplication in MoE routing --- src/maxtext/configs/base.yml | 3 +- src/maxtext/configs/types.py | 4 + src/maxtext/layers/moe.py | 240 ++++++++++++++++++++++++++++++++--- tests/unit/moe_test.py | 66 ++++++++++ 4 files changed, 295 insertions(+), 18 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 6b23e31d40..3e1a98c119 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -269,8 +269,9 @@ merge_gating_gmm: false norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights. -# when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls moe_fsdp_use_two_stage_all_gather: false +# Enable token activation deduplication in MoE ragged all-to-all communication +enable_moe_token_activation_dedup: false # Shard the expert dimension of the MLP weights on the FSDP axis. # This configuration is recommended only when num_experts is a multiple of fsdp_parallelism shard_exp_on_fsdp: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 4fe77ff057..da9d96ccf2 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -924,6 +924,10 @@ class MoEGeneral(BaseModel): False, description="Use `fsdp` and `fsdp_transpose` axes for 2D FSDP sharding.", ) + enable_moe_token_activation_dedup: bool = Field( + False, + description="Enable token activation deduplication in MoE ragged all-to-all communication.", + ) norm_topk_prob: bool = Field( False, description="Enable top-k probability normalization for router weights (Qwen3-specific).", diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 915193eb0f..d2c1afbee1 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -70,6 +70,11 @@ class RouteMetadata: # Shape [num_ep, num_ep]. all_gather of reshaped_group_sizes across EP shards. # [i, j] = number of tokens from batch shard i sent to expert shard j. all_shards_group_sizes: Optional[jax.Array] + # Token activation deduplication metadata + is_dedup: bool = False + dedup_sender_gather_indices: Optional[jax.Array] = None + dedup_receiver_expand_indices: Optional[jax.Array] = None + dedup_receiver_weights: Optional[jax.Array] = None @struct.dataclass @@ -1704,7 +1709,143 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, ), ) + def ra2a_dedup_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): + inputs_shape = x.shape + bsz_times_seq_len = inputs_shape[0] * inputs_shape[1] + inputs_2d = jnp.reshape(x, (bsz_times_seq_len, inputs_shape[2])) + weights, selected_experts = self.get_topk(logits, pre_bias_logits, rngs, input_ids) + + lb_loss = None + if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: + softmax_probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1).astype(self.dtype) + lb_loss = self.load_balance_loss(selected_experts, softmax_probs) + + if self.should_update_load_balance(): + bias_updates = calculate_load_balance_updates( + selected_experts, + self.config.num_experts, + self.config.routed_bias_update_rate, + ) + else: + bias_updates = None + + local_expert_size = self.config.num_experts // num_ep + batch_axis = self._expert_parallelism_name if is_batch_sharded_by_expert else "data" + + dest_shards = selected_experts // local_expert_size # [N, K] + shard_range = jnp.arange(num_ep)[:, None, None] # [num_ep, 1, 1] + + # token_picks: [num_ep, bsz_times_seq_len] + token_picks = jnp.any(dest_shards[None, :, :] == shard_range, axis=-1) + dedup_send_sizes = jnp.sum(token_picks, axis=-1) + + def get_shard_send_indices(d_mask): + indices = jnp.where(d_mask, jnp.arange(bsz_times_seq_len), -1) + sort_order = jnp.argsort(indices < 0) + return indices[sort_order] + + send_gather_indices = jax.vmap(get_shard_send_indices)(token_picks) # [num_ep, bsz_times_seq_len] + flat_send_indices = send_gather_indices.reshape(-1) + x_send = jnp.where(flat_send_indices[:, None] >= 0, inputs_2d[jnp.maximum(flat_send_indices, 0)], 0.0).astype(self.dtype) + + all_shards_dedup_sizes = jax.lax.all_gather(dedup_send_sizes, axis_name=batch_axis) + + buffer_size = bsz_times_seq_len * num_ep + input_offsets, send_sizes, output_offsets, recv_sizes = RoutedMoE.get_all_to_all_params( + all_shards_dedup_sizes, + expert_shard_id, + num_ep, + ragged_buffer_factor=self.config.ragged_buffer_factor, + buffer_size=buffer_size, + ) + + output_shape = jax.lax.empty((buffer_size, self.moe_expert_input_dim), dtype=x.dtype) + x_recv = jax.lax.ragged_all_to_all( + x_send, + output_shape, + input_offsets, + send_sizes, + output_offsets, + recv_sizes, + axis_name=self._expert_parallelism_name, + ) + + # All-gather routing metadata across EP axis + all_selected_experts = jax.lax.all_gather(selected_experts, axis_name=self._expert_parallelism_name) + all_weights = jax.lax.all_gather(weights, axis_name=self._expert_parallelism_name) + + # Receiver local expansion map + s_dest_shards = all_selected_experts // local_expert_size + s_picks_me = jnp.any(s_dest_shards == expert_shard_id, axis=-1) # [num_ep, N] + + zero_first = jnp.concatenate([jnp.zeros((1,), dtype=jnp.int32), jnp.cumsum(all_shards_dedup_sizes[:, expert_shard_id])[:-1]]) + s_unique_table = jax.vmap(get_shard_send_indices)(s_picks_me) + s_token_ranks = jax.vmap(lambda table: jnp.argsort(jnp.argsort(jnp.where(table >= 0, table, bsz_times_seq_len + jnp.arange(bsz_times_seq_len)))))(s_unique_table) + s_unique_slot = zero_first[:, None] + s_token_ranks # [num_ep, N] + + local_expert_range = jnp.arange(local_expert_size)[:, None, None, None] + target_exp_ids = expert_shard_id * local_expert_size + local_expert_range + match_mask = (all_selected_experts[None, :, :, :] == target_exp_ids) # [local_expert_size, num_ep, N, K] + + local_group_sizes = jnp.sum(match_mask, axis=(1, 2, 3)) + total_local_tokens = bsz_times_seq_len * self.num_experts_per_tok + + slot_broadcast = jnp.broadcast_to(s_unique_slot[None, :, :, None], match_mask.shape) + weights_broadcast = jnp.broadcast_to(all_weights[None, :, :, :], match_mask.shape) + + flat_matches = match_mask.reshape(local_expert_size, -1) + flat_slots = slot_broadcast.reshape(local_expert_size, -1) + flat_weights = weights_broadcast.reshape(local_expert_size, -1) + + def sort_expert_matches(m_row, s_row, w_row): + order = jnp.argsort(~m_row) + return jnp.where(m_row[order], s_row[order], -1), jnp.where(m_row[order], w_row[order], 0.0) + + sorted_slots_per_exp, sorted_weights_per_exp = jax.vmap(sort_expert_matches)(flat_matches, flat_slots, flat_weights) + + flat_all_exp_slots = sorted_slots_per_exp.reshape(-1) + flat_all_exp_weights = sorted_weights_per_exp.reshape(-1) + + is_valid = flat_all_exp_slots >= 0 + compact_order = jnp.argsort(~is_valid) + expand_indices = jnp.maximum(flat_all_exp_slots[compact_order][:total_local_tokens], 0) + expand_weights = flat_all_exp_weights[compact_order][:total_local_tokens] + + x_local = x_recv[expand_indices] + expert_indices = jnp.arange(local_expert_size) + sorted_experts_ids = jnp.repeat( + expert_indices, + repeats=local_group_sizes, + total_repeat_length=total_local_tokens, + ) + + return ( + x_local, + RouteOutput( + group_sizes=local_group_sizes, + selected_experts=sorted_experts_ids, + sorted_selected_experts=expand_indices, + weights=expand_weights, + lb_loss=lb_loss, + bias_updates=bias_updates, + local_group_sizes=local_group_sizes, + ), + RouteMetadata( + expert_shard_id=expert_shard_id, + local_sorted_indices=expand_indices, + all_shards_group_sizes=all_shards_dedup_sizes, + reshaped_group_sizes=dedup_send_sizes, + is_dedup=True, + dedup_sender_gather_indices=send_gather_indices, + dedup_receiver_expand_indices=expand_indices, + dedup_receiver_weights=expand_weights, + ), + ) + def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): + if self.config.enable_moe_token_activation_dedup and num_ep > 1 and is_batch_sharded_by_expert: + return ra2a_dedup_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=input_ids) + local_sorted_indices = None all_shards_group_sizes = None reshaped_group_sizes = None @@ -1958,6 +2099,32 @@ def unsort_output_and_ra2a( is_batch_sharded_by_expert, ): """Unsort tokens and return them to original shards using ragged all-to-all.""" + if route_metadata.is_dedup: + weighted_intermediate = intermediate_output * route_metadata.dedup_receiver_weights[:, None] + reduced_unique_outputs = jax.ops.segment_sum( + weighted_intermediate, + segment_ids=route_metadata.dedup_receiver_expand_indices, + num_segments=intermediate_output.shape[0], + ) + + input_offsets, send_sizes, output_offsets, recv_sizes = RoutedMoE.get_all_to_all_params( + route_metadata.all_shards_group_sizes, + route_metadata.expert_shard_id, + self.get_expert_parallelism_size(), + ragged_buffer_factor=self.config.ragged_buffer_factor, + buffer_size=intermediate_output.shape[0], + is_dispatch=False, + ) + return jax.lax.ragged_all_to_all( + reduced_unique_outputs, + output_shape, + input_offsets, + send_sizes, + output_offsets, + recv_sizes, + axis_name=self._expert_parallelism_name, + ) + if is_batch_sharded_by_expert: # locally unpermute back to the original order if self.config.use_ragged_sort: @@ -2194,24 +2361,63 @@ def _moe_body( return output, routing.lb_loss, routing.bias_updates if self.get_expert_parallelism_size() > 1: - original_inputs_first_dim = batch_size * sequence_length * self.config.num_experts_per_tok - if routing.sorted_selected_experts.shape[0] != original_inputs_first_dim: - raise ValueError("original_inputs_first_dim does not match the original tensor" " shape!") - output_shape = jax.lax.empty( - ( - original_inputs_first_dim, - self.moe_expert_input_dim // self.get_tensor_parallelism_size(), - ), - dtype=intermediate_output.dtype, - ) + if route_metadata.is_dedup: + num_ep = self.get_expert_parallelism_size() + bsz_times_seq_len = batch_size * sequence_length + return_shape = jax.lax.empty( + ( + bsz_times_seq_len * num_ep, + self.moe_expert_input_dim // self.get_tensor_parallelism_size(), + ), + dtype=intermediate_output.dtype, + ) + x_returned = unsort_output_and_ra2a( + intermediate_output, + routing, + route_metadata, + return_shape, + is_batch_sharded_by_expert, + ) + d_offsets = jnp.concatenate([ + jnp.zeros((1,), dtype=jnp.int32), + jnp.cumsum(route_metadata.all_shards_group_sizes[route_metadata.expert_shard_id, :])[:-1], + ]) + final_output = jnp.zeros( + (bsz_times_seq_len, self.moe_expert_input_dim // self.get_tensor_parallelism_size()), + dtype=self.dtype, + ) + for d in range(num_ep): + d_start = d_offsets[d] + d_count = route_metadata.reshaped_group_sizes[d] + d_returned = jax.lax.dynamic_slice_in_dim(x_returned, d_start, bsz_times_seq_len, axis=0) + t_indices = route_metadata.dedup_sender_gather_indices[d] + final_output = final_output.at[jnp.maximum(t_indices, 0)].add( + jnp.where(t_indices[:, None] >= 0, d_returned, 0.0) + ) + return ( + final_output.reshape(batch_size, sequence_length, -1).astype(self.dtype), + routing.lb_loss, + routing.bias_updates, + ) + else: + original_inputs_first_dim = batch_size * sequence_length * self.config.num_experts_per_tok + if routing.sorted_selected_experts.shape[0] != original_inputs_first_dim: + raise ValueError("original_inputs_first_dim does not match the original tensor" " shape!") + output_shape = jax.lax.empty( + ( + original_inputs_first_dim, + self.moe_expert_input_dim // self.get_tensor_parallelism_size(), + ), + dtype=intermediate_output.dtype, + ) - intermediate_output = unsort_output_and_ra2a( - intermediate_output, - routing, - route_metadata, - output_shape, - is_batch_sharded_by_expert, - ) + intermediate_output = unsort_output_and_ra2a( + intermediate_output, + routing, + route_metadata, + output_shape, + is_batch_sharded_by_expert, + ) output = self.unpermute( intermediate_output, diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 06c8713401..1d600592da 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -1054,6 +1054,72 @@ def test_megablox_expert_tensor_parallelism(self): actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh) assert_moe_close(actual_output, expected_output, cfg.dtype) + @pytest.mark.tpu_only + def test_megablox_token_activation_dedup(self): + cfg = pyconfig.initialize( + [None, get_test_config_path()], + run_name="moe_block_megablox_dedup_test", + enable_checkpointing=False, + model_name="mixtral-8x7b", + dtype="bfloat16", + megablox=True, + sparse_matmul=True, + enable_moe_token_activation_dedup=True, + per_device_batch_size=4, + ici_expert_parallelism=2, + max_target_length=128, + float32_gate_logits=True, + ) + + rng = jax.random.PRNGKey(2345) + rng_model, rng_hidden_states = jax.random.split(rng) + device_count = jax.device_count() + hidden_states = jax.random.uniform( + rng_hidden_states, + (int(cfg.per_device_batch_size) * device_count, cfg.max_target_length, cfg.base_emb_dim), + dtype=cfg.dtype, + ) + + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh) + actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh) + assert_moe_close(actual_output, expected_output, cfg.dtype) + + @pytest.mark.tpu_only + def test_tokamax_token_activation_dedup(self): + cfg = pyconfig.initialize( + [None, get_test_config_path()], + run_name="moe_block_tokamax_dedup_test", + enable_checkpointing=False, + model_name="mixtral-8x7b", + dtype="bfloat16", + use_tokamax_gmm=True, + sparse_matmul=True, + enable_moe_token_activation_dedup=True, + per_device_batch_size=4, + ici_expert_parallelism=2, + max_target_length=128, + float32_gate_logits=True, + ) + + rng = jax.random.PRNGKey(2345) + rng_model, rng_hidden_states = jax.random.split(rng) + device_count = jax.device_count() + hidden_states = jax.random.uniform( + rng_hidden_states, + (int(cfg.per_device_batch_size) * device_count, cfg.max_target_length, cfg.base_emb_dim), + dtype=cfg.dtype, + ) + + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh) + actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh) + assert_moe_close(actual_output, expected_output, cfg.dtype) + def test_random_routing(self): bs, seq_len, num_experts, num_experts_per_tok = 12, 1024, 8, 2 rng = jax.random.PRNGKey(0) From bfa71e9af6efb3c423a89cb47c6d2ab2d44dc888 Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Thu, 6 Aug 2026 14:02:57 +0000 Subject: [PATCH 2/4] Implement MoonEP-style dynamic redundant expert prefetching and rebalancing for MoE --- src/maxtext/configs/base.yml | 4 + src/maxtext/configs/types.py | 13 ++ .../synthetic_data_processing.py | 26 ++- src/maxtext/layers/moe.py | 208 +++++++++++++++++- tests/unit/moe_dynamic_redundant_e2e_test.py | 115 ++++++++++ .../moe_dynamic_redundant_planning_test.py | 127 +++++++++++ .../moe_dynamic_redundant_weights_test.py | 96 ++++++++ 7 files changed, 581 insertions(+), 8 deletions(-) create mode 100644 tests/unit/moe_dynamic_redundant_e2e_test.py create mode 100644 tests/unit/moe_dynamic_redundant_planning_test.py create mode 100644 tests/unit/moe_dynamic_redundant_weights_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 3e1a98c119..dc9db59547 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -225,6 +225,9 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss use_random_routing: false # whether to use random routing for debug/test purpose use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism +enable_moe_dynamic_redundant_experts: false # whether to use MoonEP-style dynamic redundant expert duplication and rebalancing +moe_redundant_slots_per_rank: 2 # number of prefetch slots (B) per EP rank for dynamic expert duplication +moe_rebalance_threshold_ratio: 1.15 # load threshold ratio (rank_load / avg_load) to trigger expert duplication num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe). @@ -740,6 +743,7 @@ sft_train_on_completion_only: false # dataset_type must be synthetic, hf, grain, tfds # details in: https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline.md dataset_type: tfds +synthetic_data_distribution: 'uniform' # for TFDS input pipeline (dataset_type=tfds) dataset_path: "" # your path given as argument in download_dataset.sh, e.g. "gs://my-maxtext-dataset/" dataset_name: 'c4/en:3.1.0' diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index da9d96ccf2..4f05a980fc 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -860,6 +860,18 @@ class MoEGeneral(BaseModel): False, description="Whether to use Ring of Experts for sparse matmul expert parallelism.", ) + enable_moe_dynamic_redundant_experts: bool = Field( + False, + description="Whether to use MoonEP-style dynamic redundant expert duplication and rebalancing.", + ) + moe_redundant_slots_per_rank: int = Field( + 2, + description="Number of prefetch slots (B) per EP rank for dynamic expert duplication.", + ) + moe_rebalance_threshold_ratio: float = Field( + 1.15, + description="Load threshold ratio (rank_load / avg_load) to trigger expert duplication.", + ) moe_dispatch_no_expert_sharding: bool = Field( False, description=( @@ -1340,6 +1352,7 @@ class DatasetGeneral(BaseModel): """General configuration for dataset and data loading.""" dataset_type: DatasetType = Field(DatasetType.TFDS, description="The type of the data loading pipeline.") + synthetic_data_distribution: str = Field("uniform", description="Token distribution for synthetic dataset ('uniform' or 'zipf').") per_device_batch_size: int | float = Field(12, description="The batch size per device.") eval_per_device_batch_size: int | float = Field( 0.0, diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py index d10816435c..7ec238f23d 100644 --- a/src/maxtext/input_pipeline/synthetic_data_processing.py +++ b/src/maxtext/input_pipeline/synthetic_data_processing.py @@ -41,13 +41,25 @@ def __init__(self, config, mesh): SyntheticDataIterator.raw_generate_synthetic_data, out_shardings=data_pspec_shardings, static_argnums=0 ) - tokens = jax.random.randint( - jax.random.PRNGKey(0), - (config.global_batch_size_to_load, config.max_target_length + 1), - 0, - config.vocab_size, - dtype=jnp.int32, - ) + if getattr(config, "synthetic_data_distribution", "uniform") == "zipf": + alpha = 1.25 + ranks = np.arange(1, config.vocab_size + 1, dtype=np.float64) + probs = 1.0 / (ranks ** alpha) + probs /= probs.sum() + flat_size = config.global_batch_size_to_load * (config.max_target_length + 1) + np.random.seed(42) + np_tokens = np.random.choice(config.vocab_size, size=flat_size, p=probs).reshape( + config.global_batch_size_to_load, config.max_target_length + 1 + ).astype(np.int32) + tokens = jnp.array(np_tokens) + else: + tokens = jax.random.randint( + jax.random.PRNGKey(0), + (config.global_batch_size_to_load, config.max_target_length + 1), + 0, + config.vocab_size, + dtype=jnp.int32, + ) sequence_positions = jnp.arange(0, config.max_target_length + 1, dtype=jnp.int32).reshape(1, -1) batch_positions = jnp.broadcast_to( diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index d2c1afbee1..027bcb7e53 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -117,6 +117,171 @@ def _truncate_matrix(all_shards_group_sizes: jax.Array, buffer_size: int) -> jax return jnp.diff(clamped_cumsum_extended, axis=0) +def plan_expert_duplication( + token_counts_global: jax.Array, + num_ep: int, + experts_per_rank: int, + num_slots_B: int, + rebalance_threshold: float = 1.15, +) -> tuple[jax.Array, jax.Array]: + """Pure JAX greedy load balancer for MoonEP-style dynamic redundant expert allocation. + + Args: + token_counts_global: [num_ep, E] matrix of token counts routed from each rank to each expert. + num_ep: Number of EP ranks (R). + experts_per_rank: E / R (E_local). + num_slots_B: Number of prefetch slots per rank (B). + rebalance_threshold: Minimum ratio of rank load to average load to trigger duplication. + + Returns: + slot_assignments: [num_ep, num_slots_B] int32 - Global expert ID assigned to each rank's slot (-1 if unused). + reroute_fractions: [num_experts] float32 - Fraction of tokens for each expert to divert to redundant slots. + """ + total_tokens_per_expert = jnp.sum(token_counts_global, axis=0) + num_experts = num_ep * experts_per_rank + + # 1. Compute baseline rank load: shape [num_ep] + rank_loads = jnp.sum(total_tokens_per_expert.reshape(num_ep, experts_per_rank), axis=-1).astype(jnp.float32) + avg_load = jnp.mean(rank_loads) + + # 2. State for the greedy loop: runs num_ep * num_slots_B times + max_iterations = num_ep * num_slots_B + + init_state = ( + rank_loads, + jnp.full((num_ep, num_slots_B), -1, dtype=jnp.int32), + jnp.zeros((num_ep,), dtype=jnp.int32), + jnp.zeros((num_experts,), dtype=jnp.float32), + jnp.zeros((num_experts,), dtype=jnp.bool_), + ) + + def greedy_step(i, state): + loads, slots, fill_counts, fractions, dup_mask = state + + r_max = jnp.argmax(loads) + r_min = jnp.argmin(loads) + + max_load = loads[r_max] + min_load = loads[r_min] + + has_free_slot = fill_counts[r_min] < num_slots_B + is_overloaded = max_load > (avg_load * rebalance_threshold) + can_balance = jnp.logical_and(has_free_slot, is_overloaded) + + # Find hottest un-duplicated expert belonging to r_max + r_max_expert_start = r_max * experts_per_rank + expert_indices = jnp.arange(num_experts) + r_max_expert_mask = (expert_indices >= r_max_expert_start) & (expert_indices < r_max_expert_start + experts_per_rank) + candidate_expert_scores = jnp.where( + r_max_expert_mask & (~dup_mask), + total_tokens_per_expert, + -1, + ) + e_hot = jnp.argmax(candidate_expert_scores) + hot_tokens = total_tokens_per_expert[e_hot].astype(jnp.float32) + + # Calculate tokens to divert: 50% of the hot expert's load (capped by load delta) + tokens_to_divert = jnp.minimum(hot_tokens * 0.5, (max_load - min_load) * 0.5) + divert_fraction = jnp.where(hot_tokens > 0.0, tokens_to_divert / hot_tokens, 0.0) + + # Update state conditionally + slot_idx = jnp.clip(fill_counts[r_min], 0, num_slots_B - 1) + new_slots = slots.at[r_min, slot_idx].set(jnp.where(can_balance, e_hot, slots[r_min, slot_idx])) + new_fill_counts = fill_counts.at[r_min].add(jnp.where(can_balance, 1, 0)) + new_fractions = fractions.at[e_hot].set(jnp.where(can_balance, divert_fraction, fractions[e_hot])) + new_dup_mask = dup_mask.at[e_hot].set(jnp.where(can_balance, True, dup_mask[e_hot])) + + new_loads = loads.at[r_max].add(jnp.where(can_balance, -tokens_to_divert, 0.0)) + new_loads = new_loads.at[r_min].add(jnp.where(can_balance, tokens_to_divert, 0.0)) + + return (new_loads, new_slots, new_fill_counts, new_fractions, new_dup_mask) + + final_state = jax.lax.fori_loop(0, max_iterations, greedy_step, init_state) + return final_state[1], final_state[3] + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5)) +def _manage_dynamic_expert_weights( + w_home: jax.Array, + slot_assignments: jax.Array, + num_ep: int, + experts_per_rank: int, + num_slots_B: int, + axis_name: str, +) -> jax.Array: + """Dynamic weight manager: prefetches weights into slots [B] and reduces backward slot gradients.""" + return _manage_dynamic_expert_weights_fwd( + w_home, slot_assignments, num_ep, experts_per_rank, num_slots_B, axis_name + )[0] + + +def _manage_dynamic_expert_weights_fwd( + w_home: jax.Array, + slot_assignments: jax.Array, + num_ep: int, + experts_per_rank: int, + num_slots_B: int, + axis_name: str, +) -> tuple[jax.Array, tuple[jax.Array]]: + all_weights = jax.lax.all_gather(w_home, axis_name=axis_name) + all_weights_flat = jnp.reshape(all_weights, (-1, *w_home.shape[1:])) + + my_shard_id = jax.lax.axis_index(axis_name) + my_slot_experts = slot_assignments[my_shard_id] + + valid_slot_mask = (my_slot_experts >= 0).reshape((num_slots_B,) + (1,) * (w_home.ndim - 1)) + safe_indices = jnp.maximum(my_slot_experts, 0) + prefetched_weights = jnp.take(all_weights_flat, safe_indices, axis=0) + prefetched_weights = jnp.where(valid_slot_mask, prefetched_weights, 0.0) + + w_active = jnp.concatenate([w_home, prefetched_weights], axis=0) + res = (slot_assignments,) + return w_active, res + + +def _manage_dynamic_expert_weights_bwd( + num_ep: int, + experts_per_rank: int, + num_slots_B: int, + axis_name: str, + res: tuple[jax.Array], + g_w_active: jax.Array, +) -> tuple[jax.Array, None]: + (slot_assignments,) = res + my_shard_id = jax.lax.axis_index(axis_name) + + local_e = g_w_active.shape[0] - num_slots_B + g_home = g_w_active[:local_e] + g_slots = g_w_active[local_e:] + + my_slot_experts = slot_assignments[my_shard_id] + + g_global_sparse = jnp.zeros((num_ep * local_e, *g_home.shape[1:]), dtype=g_home.dtype) + valid_slots = (my_slot_experts >= 0).reshape((num_slots_B,) + (1,) * (g_home.ndim - 1)) + safe_indices = jnp.maximum(my_slot_experts, 0) + + g_global_sparse = g_global_sparse.at[safe_indices].add( + jnp.where(valid_slots, g_slots, 0.0) + ) + + g_global_summed = jax.lax.psum(g_global_sparse, axis_name=axis_name) + + my_home_start = my_shard_id * local_e + remote_slot_grads_for_me = jax.lax.dynamic_slice_in_dim( + g_global_summed, my_home_start, local_e, axis=0 + ) + + g_w_final = g_home + remote_slot_grads_for_me + return (g_w_final, None) + + +_manage_dynamic_expert_weights.defvjp( + _manage_dynamic_expert_weights_fwd, _manage_dynamic_expert_weights_bwd +) + + + + def _sort_activations( inputs: jax.Array, sort_indices: jax.Array, @@ -2084,9 +2249,14 @@ def get_gmm_for_local_experts(x, routing, route_metadata): experts_start = route_metadata.expert_shard_id * num_experts_per_shard else: experts_start = 0 + group_sizes = routing.group_sizes + if self.config.enable_moe_dynamic_redundant_experts and num_ep > 1: + num_slots_B = self.config.moe_redundant_slots_per_rank + group_sizes = jnp.pad(group_sizes, (0, num_slots_B), constant_values=0) + return functools.partial( gmm, - group_sizes=routing.group_sizes, + group_sizes=group_sizes, expert_assignments=routing.selected_experts, group_offset=experts_start, ) @@ -2289,6 +2459,42 @@ def _moe_body( rngs, ): batch_size, sequence_length, embed_dim = x.shape + if self.config.enable_moe_dynamic_redundant_experts and self.get_expert_parallelism_size() > 1: + num_ep = self.get_expert_parallelism_size() + local_expert_size = self.config.num_experts // num_ep + num_slots_B = self.config.moe_redundant_slots_per_rank + my_shard_id = jax.lax.axis_index(self._expert_parallelism_name) + + _, selected_experts_init = self.get_topk(logits, pre_bias_logits, rngs, sharded_input_ids) + flat_selected = selected_experts_init.reshape(-1) + local_counts = jnp.bincount(flat_selected, length=self.config.num_experts) + global_counts = jax.lax.all_gather(local_counts, axis_name=self._expert_parallelism_name) + + slot_assignments, _ = plan_expert_duplication( + global_counts, + num_ep=num_ep, + experts_per_rank=local_expert_size, + num_slots_B=num_slots_B, + rebalance_threshold=self.config.moe_rebalance_threshold_ratio, + ) + + if w0.shape[0] == self.config.num_experts: + w0_home = jax.lax.dynamic_slice_in_dim(w0, my_shard_id * local_expert_size, local_expert_size, axis=0) + w1_home = jax.lax.dynamic_slice_in_dim(w1, my_shard_id * local_expert_size, local_expert_size, axis=0) + wo_home = jax.lax.dynamic_slice_in_dim(wo, my_shard_id * local_expert_size, local_expert_size, axis=0) + else: + w0_home, w1_home, wo_home = w0, w1, wo + + w0 = _manage_dynamic_expert_weights( + w0_home, slot_assignments, num_ep, local_expert_size, num_slots_B, self._expert_parallelism_name + ) + w1 = _manage_dynamic_expert_weights( + w1_home, slot_assignments, num_ep, local_expert_size, num_slots_B, self._expert_parallelism_name + ) + wo = _manage_dynamic_expert_weights( + wo_home, slot_assignments, num_ep, local_expert_size, num_slots_B, self._expert_parallelism_name + ) + if self.config.num_moe_emb_chunks > 0: output0, output1, gmm_fn, routing, route_metadata, wo_bias = moe_emb_chunking( x, diff --git a/tests/unit/moe_dynamic_redundant_e2e_test.py b/tests/unit/moe_dynamic_redundant_e2e_test.py new file mode 100644 index 0000000000..0603b155a1 --- /dev/null +++ b/tests/unit/moe_dynamic_redundant_e2e_test.py @@ -0,0 +1,115 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end forward and backward tests for MoonEP dynamic redundant expert parallelism.""" + +import unittest +import sys +from absl import flags +import jax +import jax.numpy as jnp +import numpy as np + +# Prevent absl flags crash when unittest passes CLI arguments like -s or -p +if not flags.FLAGS.is_parsed(): + try: + flags.FLAGS(sys.argv[:1]) + except Exception: + pass +from jax.sharding import Mesh +from maxtext.configs import pyconfig +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + + +class MoeDynamicRedundantE2ETest(unittest.TestCase): + """End-to-end test for dynamic redundant experts.""" + + def test_forward_and_backward_dynamic_redundancy(self): + """Executes a full forward and backward pass with dynamic redundancy enabled on TPU.""" + cfg = pyconfig.initialize( + [None, get_test_config_path()], + run_name="moe_dynamic_redundancy_e2e_test", + enable_checkpointing=False, + model_name="mixtral-8x7b", + dtype="bfloat16", + weight_dtype="bfloat16", + megablox=False, + sparse_matmul=True, + use_tokamax_gmm=True, + ici_expert_parallelism=4, + ici_fsdp_parallelism=-1, + per_device_batch_size=1, + max_target_length=128, + enable_moe_dynamic_redundant_experts=True, + moe_redundant_slots_per_rank=2, + moe_rebalance_threshold_ratio=1.15, + float32_gate_logits=True, + ) + + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + + rng = jax.random.PRNGKey(42) + rng_init, rng_data = jax.random.split(rng) + device_count = jax.device_count() + batch_size = int(cfg.per_device_batch_size) * device_count + + hidden_states = jax.random.uniform( + rng_data, + (batch_size, cfg.max_target_length, cfg.base_emb_dim), + dtype=cfg.dtype, + ) + + model = moe.get_routed_moe( + name="MoeBlock", + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + + # Initialize model variables + variables = model.init( + {"params": rng_init, "dropout": rng_init}, + hidden_states, + ) + + # Forward loss function with value_and_grad + def loss_fn(params, inputs): + output, lb_loss, bias_updates = model.apply({"params": params}, inputs) + return jnp.sum(output.astype(jnp.float32)) + + jitted_step = jax.jit(jax.value_and_grad(loss_fn)) + loss_val, grads = jitted_step(variables["params"], hidden_states) + + # Convert to numpy for assertion checks + loss_np = float(loss_val) + self.assertTrue(np.isfinite(loss_np), f"Loss is not finite: {loss_np}") + + for param_name, grad_arr in grads.items(): + if isinstance(grad_arr, jax.Array): + grad_np = np.array(grad_arr) + self.assertTrue(np.isfinite(grad_np).all(), f"Gradient for {param_name} contains NaNs or Infs!") + + print(f"End-to-end Dynamic Redundant Expert Step Successful! Loss: {loss_np:.4f}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/moe_dynamic_redundant_planning_test.py b/tests/unit/moe_dynamic_redundant_planning_test.py new file mode 100644 index 0000000000..43f15f4e2a --- /dev/null +++ b/tests/unit/moe_dynamic_redundant_planning_test.py @@ -0,0 +1,127 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for MoonEP-style dynamic redundant expert planning in MaxText.""" + +import unittest +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.layers.moe import plan_expert_duplication + + +class MoeDynamicRedundantPlanningTest(unittest.TestCase): + """Tests for the pure-JAX JIT-compiled greedy load balancer.""" + + def setUp(self): + super().setUp() + self.num_ep = 4 + self.experts_per_rank = 16 # Total experts = 64 + self.num_slots_B = 2 + self.total_experts = self.num_ep * self.experts_per_rank + + def test_balanced_distribution_no_duplication(self): + """Under balanced routing, no experts should be duplicated.""" + # Each rank routes exactly 100 tokens to each expert + token_counts = jnp.full((self.num_ep, self.total_experts), 100, dtype=jnp.int32) + + slot_assignments, reroute_fractions = plan_expert_duplication( + token_counts, + num_ep=self.num_ep, + experts_per_rank=self.experts_per_rank, + num_slots_B=self.num_slots_B, + rebalance_threshold=1.15, + ) + + # All slots should be unused (-1) + np.testing.assert_array_equal(np.array(slot_assignments), -1) + # No tokens should be rerouted + np.testing.assert_allclose(np.array(reroute_fractions), 0.0) + + def test_single_hot_expert_rebalancing(self): + """When a single expert is heavily overloaded, it should be assigned to an idle rank.""" + # Rank 0, Expert 0 gets 8000 tokens; all other experts get 50 tokens + counts = np.full((self.num_ep, self.total_experts), 50, dtype=np.int32) + counts[:, 0] = 2000 # Total for Expert 0 = 8000 tokens + token_counts = jnp.array(counts) + + slot_assignments, reroute_fractions = plan_expert_duplication( + token_counts, + num_ep=self.num_ep, + experts_per_rank=self.experts_per_rank, + num_slots_B=self.num_slots_B, + rebalance_threshold=1.15, + ) + + slots_np = np.array(slot_assignments) + fractions_np = np.array(reroute_fractions) + + # Expert 0 MUST be assigned to at least one rank's slot + self.assertIn(0, slots_np) + # The home rank for Expert 0 is Rank 0 (0 // 16 = 0). It should NOT be assigned to Rank 0's own slot + self.assertNotIn(0, slots_np[0]) + # The reroute fraction for Expert 0 should be positive (diverting load) + self.assertGreater(fractions_np[0], 0.0) + self.assertLessEqual(fractions_np[0], 0.5) + + def test_multi_expert_slot_capacity(self): + """Verifies that no rank exceeds its slot capacity B.""" + counts = np.full((self.num_ep, self.total_experts), 10, dtype=np.int32) + # Overload multiple experts on Rank 0 + counts[:, 0] = 1000 + counts[:, 1] = 1000 + counts[:, 2] = 1000 + counts[:, 3] = 1000 + token_counts = jnp.array(counts) + + slot_assignments, reroute_fractions = plan_expert_duplication( + token_counts, + num_ep=self.num_ep, + experts_per_rank=self.experts_per_rank, + num_slots_B=self.num_slots_B, + rebalance_threshold=1.15, + ) + + slots_np = np.array(slot_assignments) + # Shape must be strictly [num_ep, num_slots_B] + self.assertEqual(slots_np.shape, (self.num_ep, self.num_slots_B)) + # Each slot can hold at most 1 assigned expert + for r in range(self.num_ep): + assigned = [e for e in slots_np[r] if e >= 0] + self.assertLessEqual(len(assigned), self.num_slots_B) + + def test_jit_compilation_and_numerical_equivalence(self): + """Verifies JIT compilation succeeds with static shapes and matches un-jitted output.""" + counts = np.random.RandomState(42).randint(10, 500, size=(self.num_ep, self.total_experts)).astype(np.int32) + counts[:, 5] = 4000 # Make expert 5 hot + token_counts = jnp.array(counts) + + jitted_plan = jax.jit( + plan_expert_duplication, + static_argnames=("num_ep", "experts_per_rank", "num_slots_B"), + ) + + slots_unjitted, fracs_unjitted = plan_expert_duplication( + token_counts, self.num_ep, self.experts_per_rank, self.num_slots_B + ) + slots_jitted, fracs_jitted = jitted_plan( + token_counts, self.num_ep, self.experts_per_rank, self.num_slots_B + ) + + np.testing.assert_array_equal(np.array(slots_jitted), np.array(slots_unjitted)) + np.testing.assert_allclose(np.array(fracs_jitted), np.array(fracs_unjitted), rtol=1e-5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/moe_dynamic_redundant_weights_test.py b/tests/unit/moe_dynamic_redundant_weights_test.py new file mode 100644 index 0000000000..3d16ff4dec --- /dev/null +++ b/tests/unit/moe_dynamic_redundant_weights_test.py @@ -0,0 +1,96 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for MoonEP dynamic expert weight prefetch & backward gradient reducer in MaxText.""" + +import unittest +import functools +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, PartitionSpec as P + +from maxtext.layers.moe import _manage_dynamic_expert_weights + + +class MoeDynamicRedundantWeightsTest(unittest.TestCase): + """Tests for the dynamic weight manager custom VJP.""" + + def setUp(self): + super().setUp() + self.devices = jax.devices() + self.num_devices = len(self.devices) + self.num_ep = self.num_devices + self.experts_per_rank = 4 + self.num_slots_B = 2 + self.H = 16 + self.H_ffn = 32 + self.mesh = Mesh(np.array(self.devices), ("expert",)) + + def test_forward_prefetch_and_backward_grad_reduction(self): + """Verifies forward prefetch copies correct weights and backward correctly sums slot gradients.""" + num_total_experts = self.num_ep * self.experts_per_rank + + # Initialize global weights: [num_ep, experts_per_rank, H, H_ffn] + np_w = (np.arange(num_total_experts * self.H * self.H_ffn, dtype=np.float32) + 1.0).reshape( + self.num_ep, self.experts_per_rank, self.H, self.H_ffn + ) + w_sharded = jax.device_put(jnp.array(np_w), jax.sharding.NamedSharding(self.mesh, P("expert", None, None, None))) + + # Plan: duplicate Expert 0 (owned by Rank 0) to Rank 1 Slot 0 + # and duplicate Expert (experts_per_rank + 1) (owned by Rank 1) to Rank 0 Slot 0 + slot_assignments_np = np.full((self.num_ep, self.num_slots_B), -1, dtype=np.int32) + if self.num_ep > 1: + slot_assignments_np[1, 0] = 0 # Rank 1 gets Expert 0 in slot 0 + slot_assignments_np[0, 0] = self.experts_per_rank + 1 # Rank 0 gets an expert from Rank 1 + slot_assignments = jnp.array(slot_assignments_np) + + num_ep = self.num_ep + experts_per_rank = self.experts_per_rank + num_slots_B = self.num_slots_B + + # Forward loss function under shard_map + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=(P("expert", None, None, None), P()), + out_specs=P(), + check_vma=False, + ) + def compute_loss(w_local, slots): + # w_local shape inside shard_map: [1, experts_per_rank, H, H_ffn] -> squeeze out the leading shard dim + w_local_squeezed = jnp.squeeze(w_local, axis=0) + w_active = _manage_dynamic_expert_weights( + w_local_squeezed, slots, num_ep, experts_per_rank, num_slots_B, "expert" + ) + # Compute dummy loss: sum of squares + local_sum = jnp.sum(w_active * w_active) + return jax.lax.psum(local_sum, axis_name="expert") + + loss, grads = jax.value_and_grad(compute_loss)(w_sharded, slot_assignments) + + # Convert grads to numpy + grads_np = np.array(grads) + + # Mathematical Verification: + # d(sum(w_active^2))/dw_home[e] should equal 2 * w[e] * (1 + number_of_times_duplicated) + expected_grads = 2.0 * np_w + if self.num_ep > 1: + expected_grads[0, 0] = 4.0 * np_w[0, 0] # Expert 0 was duplicated to Rank 1 + expected_grads[1, 1] = 4.0 * np_w[1, 1] # Expert (experts_per_rank + 1) was duplicated to Rank 0 + + np.testing.assert_allclose(grads_np, expected_grads, rtol=1e-5, atol=1e-5) + + +if __name__ == "__main__": + unittest.main() From f5ad73150b8d56fff04f80d397839ef80d5020bc Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Thu, 6 Aug 2026 14:26:22 +0000 Subject: [PATCH 3/4] Optimize MoonEP dynamic expert weights with fused targeted all_to_all prefetching --- src/maxtext/layers/moe.py | 134 ++++++++++++------ tests/unit/moe_dynamic_redundant_e2e_test.py | 4 +- .../moe_dynamic_redundant_weights_test.py | 4 +- 3 files changed, 96 insertions(+), 46 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 027bcb7e53..18308d339b 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -201,7 +201,7 @@ def greedy_step(i, state): @functools.partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5)) -def _manage_dynamic_expert_weights( +def _manage_dynamic_expert_weights_targeted( w_home: jax.Array, slot_assignments: jax.Array, num_ep: int, @@ -209,77 +209,126 @@ def _manage_dynamic_expert_weights( num_slots_B: int, axis_name: str, ) -> jax.Array: - """Dynamic weight manager: prefetches weights into slots [B] and reduces backward slot gradients.""" - return _manage_dynamic_expert_weights_fwd( + """Targeted dynamic expert weight prefetcher: uses all_to_all to fetch only requested redundant experts.""" + return _manage_dynamic_expert_weights_targeted_fwd( w_home, slot_assignments, num_ep, experts_per_rank, num_slots_B, axis_name )[0] -def _manage_dynamic_expert_weights_fwd( +def _manage_dynamic_expert_weights_targeted_fwd( w_home: jax.Array, slot_assignments: jax.Array, num_ep: int, experts_per_rank: int, num_slots_B: int, axis_name: str, -) -> tuple[jax.Array, tuple[jax.Array]]: - all_weights = jax.lax.all_gather(w_home, axis_name=axis_name) - all_weights_flat = jnp.reshape(all_weights, (-1, *w_home.shape[1:])) - +): my_shard_id = jax.lax.axis_index(axis_name) - my_slot_experts = slot_assignments[my_shard_id] + requested_experts = slot_assignments + owner_ranks = jnp.where(requested_experts >= 0, requested_experts // experts_per_rank, -1) + local_indices = jnp.where(requested_experts >= 0, requested_experts % experts_per_rank, 0) + + is_owner = (owner_ranks == my_shard_id) + safe_local_indices = jnp.maximum(local_indices, 0) + gathered_local = jnp.take(w_home, safe_local_indices, axis=0) + + mask = is_owner.reshape((num_ep, num_slots_B) + (1,) * (w_home.ndim - 1)) + to_send = jnp.where(mask, gathered_local, 0.0) - valid_slot_mask = (my_slot_experts >= 0).reshape((num_slots_B,) + (1,) * (w_home.ndim - 1)) - safe_indices = jnp.maximum(my_slot_experts, 0) - prefetched_weights = jnp.take(all_weights_flat, safe_indices, axis=0) - prefetched_weights = jnp.where(valid_slot_mask, prefetched_weights, 0.0) + received = jax.lax.all_to_all(to_send, axis_name=axis_name, split_axis=0, concat_axis=0) + prefetched_weights = jnp.sum(received, axis=0) w_active = jnp.concatenate([w_home, prefetched_weights], axis=0) - res = (slot_assignments,) + res = (slot_assignments, owner_ranks, local_indices, is_owner) return w_active, res -def _manage_dynamic_expert_weights_bwd( +def _manage_dynamic_expert_weights_targeted_bwd( num_ep: int, experts_per_rank: int, num_slots_B: int, axis_name: str, - res: tuple[jax.Array], + res: tuple, g_w_active: jax.Array, -) -> tuple[jax.Array, None]: - (slot_assignments,) = res +): + slot_assignments, owner_ranks, local_indices, is_owner = res my_shard_id = jax.lax.axis_index(axis_name) - local_e = g_w_active.shape[0] - num_slots_B + local_e = experts_per_rank g_home = g_w_active[:local_e] g_slots = g_w_active[local_e:] - my_slot_experts = slot_assignments[my_shard_id] - - g_global_sparse = jnp.zeros((num_ep * local_e, *g_home.shape[1:]), dtype=g_home.dtype) - valid_slots = (my_slot_experts >= 0).reshape((num_slots_B,) + (1,) * (g_home.ndim - 1)) - safe_indices = jnp.maximum(my_slot_experts, 0) + my_requested_owners = owner_ranks[my_shard_id] + rank_indices = jnp.arange(num_ep) + target_mask = (rank_indices.reshape(num_ep, 1) == my_requested_owners.reshape(1, num_slots_B)) + mask = target_mask.reshape((num_ep, num_slots_B) + (1,) * (g_home.ndim - 1)) - g_global_sparse = g_global_sparse.at[safe_indices].add( - jnp.where(valid_slots, g_slots, 0.0) + g_slots_expanded = jnp.broadcast_to( + g_slots.reshape((1, num_slots_B) + g_slots.shape[1:]), + (num_ep, num_slots_B) + g_slots.shape[1:], ) + g_to_send = jnp.where(mask, g_slots_expanded, 0.0) - g_global_summed = jax.lax.psum(g_global_sparse, axis_name=axis_name) + g_received = jax.lax.all_to_all(g_to_send, axis_name=axis_name, split_axis=0, concat_axis=0) - my_home_start = my_shard_id * local_e - remote_slot_grads_for_me = jax.lax.dynamic_slice_in_dim( - g_global_summed, my_home_start, local_e, axis=0 - ) + flat_is_owner = is_owner.reshape(-1) + flat_local_idx = local_indices.reshape(-1) + flat_g_rec = g_received.reshape((-1,) + g_home.shape[1:]) - g_w_final = g_home + remote_slot_grads_for_me + safe_local_idx = jnp.maximum(flat_local_idx, 0) + mask_rec = flat_is_owner.reshape((-1,) + (1,) * (g_home.ndim - 1)) + valid_g_rec = jnp.where(mask_rec, flat_g_rec, 0.0) + + g_home_accum = jnp.zeros_like(g_home).at[safe_local_idx].add(valid_g_rec) + + g_w_final = g_home + g_home_accum return (g_w_final, None) -_manage_dynamic_expert_weights.defvjp( - _manage_dynamic_expert_weights_fwd, _manage_dynamic_expert_weights_bwd +_manage_dynamic_expert_weights_targeted.defvjp( + _manage_dynamic_expert_weights_targeted_fwd, _manage_dynamic_expert_weights_targeted_bwd ) +def manage_moe_layer_weights_fused( + w0_home: jax.Array, + w1_home: jax.Array, + wo_home: jax.Array, + slot_assignments: jax.Array, + num_ep: int, + experts_per_rank: int, + num_slots_B: int, + axis_name: str, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Fused prefetching for w0, w1, and wo using a single targeted all_to_all call per layer.""" + local_e = experts_per_rank + dim0 = int(np.prod(w0_home.shape[1:])) + dim1 = int(np.prod(w1_home.shape[1:])) + dimo = int(np.prod(wo_home.shape[1:])) + + w0_flat = jnp.reshape(w0_home, (local_e, dim0)) + w1_flat = jnp.reshape(w1_home, (local_e, dim1)) + wo_flat = jnp.reshape(wo_home, (local_e, dimo)) + + w_fused = jnp.concatenate([w0_flat, w1_flat, wo_flat], axis=1) + + w_fused_active = _manage_dynamic_expert_weights_targeted( + w_fused, slot_assignments, num_ep, experts_per_rank, num_slots_B, axis_name + ) + + total_active_e = local_e + num_slots_B + w0_active_flat = w_fused_active[:, :dim0] + w1_active_flat = w_fused_active[:, dim0:dim0 + dim1] + wo_active_flat = w_fused_active[:, dim0 + dim1:] + + w0_active = jnp.reshape(w0_active_flat, (total_active_e, *w0_home.shape[1:])) + w1_active = jnp.reshape(w1_active_flat, (total_active_e, *w1_home.shape[1:])) + wo_active = jnp.reshape(wo_active_flat, (total_active_e, *wo_home.shape[1:])) + + return w0_active, w1_active, wo_active + + + def _sort_activations( @@ -2485,14 +2534,15 @@ def _moe_body( else: w0_home, w1_home, wo_home = w0, w1, wo - w0 = _manage_dynamic_expert_weights( - w0_home, slot_assignments, num_ep, local_expert_size, num_slots_B, self._expert_parallelism_name - ) - w1 = _manage_dynamic_expert_weights( - w1_home, slot_assignments, num_ep, local_expert_size, num_slots_B, self._expert_parallelism_name - ) - wo = _manage_dynamic_expert_weights( - wo_home, slot_assignments, num_ep, local_expert_size, num_slots_B, self._expert_parallelism_name + w0, w1, wo = manage_moe_layer_weights_fused( + w0_home, + w1_home, + wo_home, + slot_assignments, + num_ep, + local_expert_size, + num_slots_B, + self._expert_parallelism_name, ) if self.config.num_moe_emb_chunks > 0: diff --git a/tests/unit/moe_dynamic_redundant_e2e_test.py b/tests/unit/moe_dynamic_redundant_e2e_test.py index 0603b155a1..566d972e8e 100644 --- a/tests/unit/moe_dynamic_redundant_e2e_test.py +++ b/tests/unit/moe_dynamic_redundant_e2e_test.py @@ -49,8 +49,8 @@ def test_forward_and_backward_dynamic_redundancy(self): megablox=False, sparse_matmul=True, use_tokamax_gmm=True, - ici_expert_parallelism=4, - ici_fsdp_parallelism=-1, + ici_expert_parallelism=len(jax.devices()), + ici_fsdp_parallelism=1, per_device_batch_size=1, max_target_length=128, enable_moe_dynamic_redundant_experts=True, diff --git a/tests/unit/moe_dynamic_redundant_weights_test.py b/tests/unit/moe_dynamic_redundant_weights_test.py index 3d16ff4dec..6a28f3bab5 100644 --- a/tests/unit/moe_dynamic_redundant_weights_test.py +++ b/tests/unit/moe_dynamic_redundant_weights_test.py @@ -20,7 +20,7 @@ import numpy as np from jax.sharding import Mesh, PartitionSpec as P -from maxtext.layers.moe import _manage_dynamic_expert_weights +from maxtext.layers.moe import _manage_dynamic_expert_weights_targeted class MoeDynamicRedundantWeightsTest(unittest.TestCase): @@ -70,7 +70,7 @@ def test_forward_prefetch_and_backward_grad_reduction(self): def compute_loss(w_local, slots): # w_local shape inside shard_map: [1, experts_per_rank, H, H_ffn] -> squeeze out the leading shard dim w_local_squeezed = jnp.squeeze(w_local, axis=0) - w_active = _manage_dynamic_expert_weights( + w_active = _manage_dynamic_expert_weights_targeted( w_local_squeezed, slots, num_ep, experts_per_rank, num_slots_B, "expert" ) # Compute dummy loss: sum of squares From dae8fb83772c8ec157dc5fc56fa1629f7d322454 Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Thu, 6 Aug 2026 15:46:14 +0000 Subject: [PATCH 4/4] Optimize default prefetch slots B=1 for 50% reduced MoonEP weight overhead --- src/maxtext/configs/base.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index dc9db59547..ac52493b2e 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -226,7 +226,7 @@ use_random_routing: false # whether to use random routing for debug/test purpose use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism enable_moe_dynamic_redundant_experts: false # whether to use MoonEP-style dynamic redundant expert duplication and rebalancing -moe_redundant_slots_per_rank: 2 # number of prefetch slots (B) per EP rank for dynamic expert duplication +moe_redundant_slots_per_rank: 1 # number of prefetch slots (B) per EP rank for dynamic expert duplication moe_rebalance_threshold_ratio: 1.15 # load threshold ratio (rank_load / avg_load) to trigger expert duplication num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM