From 8c3ce92bc83d2d079a81c0d09294629cb4ef4879 Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Tue, 11 Aug 2026 01:32:55 +0000 Subject: [PATCH] Support inserting and replaying router expert decisions in trainer - In train.py loss_fn, extract forced_routed_experts from batch data and pass to model forward - In RoutedMoE, override top_k_indices and route tokens with forced_routed_experts across sparse_matmul and dense_matmul - In Decoder and NNXDecoder, plumb forced_routed_experts down across scanned/sequential decoder layers - Plumbed forced_routed_experts across model implementations (Qwen3, Qwen3.5, DeepSeek, Gemma4, Mixtral) - Added unit tests in tests/unit/forced_routing_test.py - Added integration test in tests/test_trainer_router_replay.py verifying loss computation with forced routing --- .../integration/tunix/tunix_adapter.py | 2 + src/maxtext/layers/decoders.py | 49 +++++- src/maxtext/layers/moe.py | 165 +++++++++++++----- src/maxtext/layers/nnx_decoders.py | 49 +++++- src/maxtext/layers/nnx_wrappers.py | 20 +++ src/maxtext/models/deepseek.py | 13 +- src/maxtext/models/gemma4.py | 8 +- src/maxtext/models/mixtral.py | 3 +- src/maxtext/models/models.py | 6 + src/maxtext/models/qwen3.py | 14 +- src/maxtext/models/qwen3_5.py | 5 +- src/maxtext/trainers/pre_train/train.py | 2 + tests/test_trainer_router_replay.py | 126 +++++++++++++ tests/unit/forced_routing_test.py | 109 ++++++++++++ 14 files changed, 503 insertions(+), 68 deletions(-) create mode 100644 tests/test_trainer_router_replay.py create mode 100644 tests/unit/forced_routing_test.py diff --git a/src/maxtext/integration/tunix/tunix_adapter.py b/src/maxtext/integration/tunix/tunix_adapter.py index a5bcbf62b2..392e746ef3 100644 --- a/src/maxtext/integration/tunix/tunix_adapter.py +++ b/src/maxtext/integration/tunix/tunix_adapter.py @@ -108,6 +108,7 @@ def __call__( attention_mask: Optional[Array], # [B, L, L] or None decoder_segment_ids: Optional[Array] = None, output_hidden_states: bool = False, # ignored + forced_routed_experts: Optional[Array] = None, ) -> Tuple[Array, None]: """Forward compatible with Tunix Trainers default loss. Returns logits, None. @@ -118,6 +119,7 @@ def __call__( decoder_input_tokens=input_tokens, decoder_positions=positions, decoder_segment_ids=decoder_segment_ids, + forced_routed_experts=forced_routed_experts, ) return logits, None diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..ad69d33356 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -861,11 +861,15 @@ def __call__( kv_caches: list[jax.Array] | None = None, attention_metadata=None, deepstack_visual_embeds: None | list[jnp.ndarray] = None, + forced_routed_experts: jnp.ndarray | None = None, ): cfg = self.config mesh = self.mesh assert decoder_input_tokens.ndim == 2 # [batch, len] + if cfg.scan_layers and forced_routed_experts is not None: + raise NotImplementedError("Forced routing with scanned layers is not supported yet.") + # [batch, length] -> [batch, length, emb_dim] y = self._apply_embedding( shared_embedding, @@ -1164,6 +1168,10 @@ def __call__( global_layer_idx = global_layer_idx_offset + index kv_cache = kv_caches[index] if kv_caches is not None else None input_tokens = decoder_input_tokens if cfg.engram_layers else None + current_forced_routed_experts = None + if forced_routed_experts is not None and layer_prefix == "moe_layers": + current_forced_routed_experts = forced_routed_experts[:, :, index, :] + y, kv_cache = layer( config=cfg, mesh=mesh, @@ -1182,6 +1190,7 @@ def __call__( kv_cache=kv_cache, attention_metadata=attention_metadata, decoder_input_tokens=input_tokens, + forced_routed_experts=current_forced_routed_experts, ) if kv_caches is not None and kv_cache is not None: kv_caches[index] = kv_cache @@ -1201,6 +1210,7 @@ def __call__( slot=slot, ) else: + moe_lyr_idx = 0 for lyr in range(cfg.num_decoder_layers): RemattedBlockLayer = RemattedBlockLayers[0] layer_kwargs = {} @@ -1244,17 +1254,46 @@ def __call__( layer = RemattedBlockLayer( config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **layer_kwargs ) + current_forced_routed_experts = None + is_moe = False + if cfg.decoder_block in ( + DecoderBlockType.MIXTRAL, + DecoderBlockType.QWEN3_MOE, + DecoderBlockType.QWEN3_NEXT, + DecoderBlockType.QWEN3_5, + DecoderBlockType.QWEN3_CUSTOM_MOE, + ): + is_moe = True + elif cfg.decoder_block == DecoderBlockType.LLAMA4: + is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step) + + if is_moe and forced_routed_experts is not None: + if forced_routed_experts.ndim == 4: + current_forced_routed_experts = forced_routed_experts[:, :, moe_lyr_idx, :] + else: + current_forced_routed_experts = forced_routed_experts + moe_lyr_idx += 1 + elif is_moe: + moe_lyr_idx += 1 + + call_kwargs = { + "previous_chunk": previous_chunk, + "slot": slot, + "kv_cache": kv_cache, + "attention_metadata": attention_metadata, + } + call_kwargs.update(layer_call_kwargs) + + if is_moe and current_forced_routed_experts is not None: + call_kwargs["forced_routed_experts"] = current_forced_routed_experts + y, returned_cache = layer( y, decoder_segment_ids, decoder_positions, deterministic, model_mode, - previous_chunk=previous_chunk, - slot=slot, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - **layer_call_kwargs, + **call_kwargs, ) if kv_caches is not None and returned_cache is not None: kv_caches[lyr] = returned_cache diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f6e89f204d..4fdeabbf61 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -705,37 +705,47 @@ def should_update_load_balance(self): """ return self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and not self.is_hash_routing - def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): + def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None, forced_routed_experts=None): """get topk.""" # shape of top_k_weights & top_k_indices: # (batch, sequence, num_experts_per_tok). - if self.config.use_random_routing: - if rngs is None: - raise ValueError("The random key cannot be None for random routing.") - # Reuse the 'params' RNG stream to ensure random routing - rng = rngs.params() if hasattr(rngs, "params") and callable(getattr(rngs, "params")) else rngs - top_k_weights, top_k_indices = random_routing(rng, gate_logits, self.num_experts_per_tok) - return top_k_weights, top_k_indices - - if self.is_hash_routing: - if input_ids is None: - raise ValueError("input_ids cannot be None when is_hash_routing is True") - # Access the static routing table - tid2eid_int = self.tid2eid.value - # Cast the float32 array to int32 (JAX automatically assigns 0.0 gradients to integer casts) - tid2eid_int = tid2eid_int.astype(jnp.int32) - # Cast input_ids to int32 to safely index the hash routing table - top_k_indices = tid2eid_int[input_ids.astype(jnp.int32)] - top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) - # NOTE: deepseek2 has a different pattern - elif self.config.model_name.startswith(("deepseek3", "deepseek4")): - top_k_weights, top_k_indices = self.deepseek_routing(gate_logits, pre_bias_logits) - elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: - router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) - _, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) - top_k_weights = jnp.take_along_axis(router_probs, top_k_indices, axis=-1).astype(self.dtype) + if forced_routed_experts is not None: + top_k_indices = forced_routed_experts + if self.config.model_name.startswith(("deepseek3", "deepseek4")): + top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) + elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: + router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) + top_k_weights = jnp.take_along_axis(router_probs, top_k_indices, axis=-1).astype(self.dtype) + else: + top_k_weights = jnp.take_along_axis(gate_logits, top_k_indices, axis=-1) else: - top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) + if self.config.use_random_routing: + if rngs is None: + raise ValueError("The random key cannot be None for random routing.") + # Reuse the 'params' RNG stream to ensure random routing + rng = rngs.params() if hasattr(rngs, "params") and callable(getattr(rngs, "params")) else rngs + top_k_weights, top_k_indices = random_routing(rng, gate_logits, self.num_experts_per_tok) + return top_k_weights, top_k_indices + + if self.is_hash_routing: + if input_ids is None: + raise ValueError("input_ids cannot be None when is_hash_routing is True") + # Access the static routing table + tid2eid_int = self.tid2eid.value + # Cast the float32 array to int32 (JAX automatically assigns 0.0 gradients to integer casts) + tid2eid_int = tid2eid_int.astype(jnp.int32) + # Cast input_ids to int32 to safely index the hash routing table + top_k_indices = tid2eid_int[input_ids.astype(jnp.int32)] + top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) + # NOTE: deepseek2 has a different pattern + elif self.config.model_name.startswith(("deepseek3", "deepseek4")): + top_k_weights, top_k_indices = self.deepseek_routing(gate_logits, pre_bias_logits) + elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: + router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) + _, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) + top_k_weights = jnp.take_along_axis(router_probs, top_k_indices, axis=-1).astype(self.dtype) + else: + top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4): top_k_weights = self.deepseek_scale_weights(top_k_weights) @@ -743,9 +753,14 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): if self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): top_k_weights = jax.nn.softmax(top_k_weights.astype(jnp.float32), axis=-1).astype(self.dtype) - # Normalization of router weights (e.g. used by Qwen3, Gemma4). - if self.config.norm_topk_prob: - top_k_weights /= top_k_weights.sum(axis=-1, keepdims=True) + # Zero out weights for padding indices! + if forced_routed_experts is not None: + valid_token_mask = top_k_indices[:, :, 0] != -1 + top_k_weights = top_k_weights * valid_token_mask[:, :, None] + + # Normalization of router weights (e.g. used by Qwen3, Gemma4). + if self.config.norm_topk_prob: + top_k_weights /= top_k_weights.sum(axis=-1, keepdims=True) return top_k_weights, top_k_indices @@ -863,13 +878,15 @@ def permute( rngs=None, roll_to_expert_id=None, input_ids=None, + forced_routed_experts=None, ): """Permute tokens to group by expert to fit gmm call.""" # reshape inputs (batch, sequence, emb) to (batch * sequence, emb) inputs_shape = inputs.shape bsz_times_seq_len = inputs_shape[0] * inputs_shape[1] inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2])) - weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids) + weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids, forced_routed_experts) + lb_loss = None if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: softmax_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1).astype(self.dtype) @@ -936,13 +953,22 @@ def permute( if roll_to_expert_id is not None: flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % self.num_experts - sorted_selected_experts = jnp.argsort(flatten_selected_experts) + + if forced_routed_experts is not None: + # Fix padding bug: map -1 to dummy valid indices to distribute load + valid_mask = flatten_selected_experts >= 0 + dummy_indices = jnp.arange(flatten_selected_experts.shape[0]) % self.num_experts + flatten_selected_experts_safe = jnp.where(valid_mask, flatten_selected_experts, dummy_indices) + else: + flatten_selected_experts_safe = flatten_selected_experts + + sorted_selected_experts = jnp.argsort(flatten_selected_experts_safe) # sort inputs for number of selected experts replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0) sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype( self.dtype ) - group_size = jnp.bincount(flatten_selected_experts, length=self.num_experts) + group_size = jnp.bincount(flatten_selected_experts_safe, length=self.num_experts) num_tokens = bsz_times_seq_len * self.num_experts_per_tok use_truncated_buffer = use_ragged_in_permute and buffer_size is not None and buffer_size < num_tokens @@ -974,7 +1000,6 @@ def permute( repeats=group_size, total_repeat_length=math.prod(selected_experts.shape), ) - return ( sorted_inputs, sorted_selected_experts, @@ -1397,6 +1422,7 @@ def sparse_matmul( w1_bias, wo_bias, input_ids=None, + forced_routed_experts=None, ): """Perform sparse matrix multiplication of inputs and Experts.""" @@ -1659,7 +1685,7 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): ) = get_routed_moe_shardings(is_batch_sharded_by_expert, input_ids is not None) w0_pspec, w1_pspec, wo_pspec = maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec) - def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): + def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None, forced_routed_experts=None): # The ring-of-experts strategy first duplicates the inputs to all # expert shards, and then routes within each shard. @@ -1687,6 +1713,7 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, roll_to_expert_id=num_experts_per_shard * expert_shard_id, rngs=rngs, input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) return ( x, @@ -1707,7 +1734,7 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, ), ) - def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): + def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None, forced_routed_experts=None): local_sorted_indices = None all_shards_group_sizes = None reshaped_group_sizes = None @@ -1727,6 +1754,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in self.config.use_custom_sort_vjp, rngs, input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) if num_ep > 1: @@ -1808,7 +1836,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in ), ) - def route(x, logits, pre_bias_logits, rngs, input_ids=None): + def route(x, logits, pre_bias_logits, rngs, input_ids=None, forced_routed_experts=None): """Performs both across device and within device token routing/sorting""" num_ep = self.get_expert_parallelism_size() expert_shard_id = jax.lax.axis_index(self._expert_parallelism_name) if num_ep > 1 else 0 @@ -1822,6 +1850,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None): expert_shard_id, rngs, input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) else: return ra2a_and_route( @@ -1832,6 +1861,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None): expert_shard_id, rngs, input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index): @@ -2123,6 +2153,7 @@ def _moe_body( wo_bias, sharded_input_ids, rngs, + forced_routed_experts=None, ): batch_size, sequence_length, embed_dim = x.shape if self.config.num_moe_emb_chunks > 0: @@ -2140,7 +2171,9 @@ def _moe_body( embed_dim, ) else: - x, routing, route_metadata = route(x, logits, pre_bias_logits, rngs, input_ids=sharded_input_ids) + x, routing, route_metadata = route( + x, logits, pre_bias_logits, rngs, input_ids=sharded_input_ids, forced_routed_experts=forced_routed_experts + ) if self.config.mlp_bias: w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias) @@ -2243,6 +2276,9 @@ def _moe_body( wo_bias_pspec, decoder_tokens_pspec, P(), # Replicate the input key + self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) + if forced_routed_experts is not None + else None, ), out_specs=( self._logical_to_mesh_axes( @@ -2269,6 +2305,7 @@ def sparse_matmul_route_and_compute( wo_bias, sharded_input_ids, rngs, + forced_routed_experts=None, ): # The expert weights (w0/w1/wo) are all-gathered over FSDP once at this # shard_map entry (implicitly, via the `embed_tensor_transpose` pspec which @@ -2288,6 +2325,7 @@ def sparse_matmul_route_and_compute( wo_bias, sharded_input_ids, rngs, + forced_routed_experts, ) # Chunked ring-of-experts pipeline: split the per-shard tokens along the @@ -2322,6 +2360,7 @@ def sparse_matmul_route_and_compute( wo_bias, None if sharded_input_ids is None else sharded_input_ids[:, sl], rngs, + None if forced_routed_experts is None else forced_routed_experts[:, sl, :], ) if self.config.moe_chunk_barrier: _prev = out_c @@ -2364,6 +2403,11 @@ def sparse_matmul_route_and_compute( inputs = self._maybe_shard_with_logical(inputs, input_axes) gate_logits = self._maybe_shard_with_logical(gate_logits, gate_logits_axes) pre_bias_logits = self._maybe_shard_with_logical(pre_bias_logits, pre_bias_logits_axes) + forced_routed_experts = ( + self._maybe_shard_with_logical(forced_routed_experts, (batch_logical_axis, "activation_norm_length", None)) + if forced_routed_experts is not None + else None + ) w0_kernel = self._maybe_shard_with_pspec(w0_kernel, w0_pspec) w1_kernel = self._maybe_shard_with_pspec(w1_kernel, w1_pspec) @@ -2387,27 +2431,36 @@ def sparse_matmul_route_and_compute( wo_bias, input_ids, self.rngs, + forced_routed_experts, ) - def reshape_and_update_weights(self, weights, indices): + def reshape_and_update_weights(self, weights, indices, safe_updates=False): """reshape and update weights.""" # input of weights and indices: (batch_size, seq_len, num_experts_per_tok) # output of updated weights: (batch_size, seq_len, num_experts) update_weights = jnp.zeros((weights.shape[0], weights.shape[1], self.num_experts), dtype=self.dtype) + if safe_updates: + valid_mask = indices >= 0 + safe_indices = jnp.where(valid_mask, indices, 0) + safe_weights = jnp.where(valid_mask, weights, 0.0) + else: + safe_indices = indices + safe_weights = weights + index_update = ( self._maybe_shard_with_logical( jnp.arange(weights.shape[0])[:, None, None], ("activation_batch", None, None), ), self._maybe_shard_with_logical(jnp.arange(weights.shape[1])[:, None], ("activation_length", None)), - indices, + safe_indices, ) weight_sharding = ( create_sharding(self.mesh, ("activation_batch", "activation_length", None)) if self.config.shard_mode == ShardMode.EXPLICIT else None ) - update_weights = update_weights.at[index_update].set(weights, out_sharding=weight_sharding) + update_weights = update_weights.at[index_update].set(safe_weights, out_sharding=weight_sharding) return update_weights def get_context_partition_and_sub_seq(self, seq_len): @@ -2666,6 +2719,7 @@ def dense_matmul( w1_bias, wo_bias, input_ids=None, + forced_routed_experts=None, ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]: """Dense matrix multiplication.""" # gate_logits: batch, length, expert @@ -2676,13 +2730,21 @@ def dense_matmul( pre_bias_logits = self._maybe_shard_with_logical( pre_bias_logits, ("activation_batch_moe", "activation_length_moe", None) ) - top_k_weights, top_k_indices = self.get_topk(gate_logits, pre_bias_logits, self.rngs, input_ids=input_ids) + if forced_routed_experts is not None: + forced_routed_experts = self._maybe_shard_with_logical( + forced_routed_experts, ("activation_batch_moe", "activation_length_moe", None) + ) + top_k_weights, top_k_indices = self.get_topk( + gate_logits, pre_bias_logits, self.rngs, input_ids=input_ids, forced_routed_experts=forced_routed_experts + ) is_llama4_decoder_layer = self.config.decoder_block == ctypes.DecoderBlockType.LLAMA4 if is_llama4_decoder_layer: router_scores = jax.nn.sigmoid(top_k_weights.astype(jnp.float32)).astype(self.dtype) inputs = inputs * router_scores else: - weights = self.reshape_and_update_weights(top_k_weights, top_k_indices) + weights = self.reshape_and_update_weights( + top_k_weights, top_k_indices, safe_updates=(forced_routed_experts is not None) + ) matmul_precision = jax.lax.Precision(self.config.matmul_precision) # Calculate load balance loss @@ -2951,7 +3013,9 @@ def dense_matmul( intermediate_layer = adc.checkpoint_name(adc.checkpoint_name(intermediate_layer, "mlpwo"), "moe_mlpwo") with jax.named_scope("weight_sum"): if is_llama4_decoder_layer: - weights = self.reshape_and_update_weights(jnp.ones_like(top_k_weights), top_k_indices) + weights = self.reshape_and_update_weights( + jnp.ones_like(top_k_weights), top_k_indices, safe_updates=(forced_routed_experts is not None) + ) if self.config.float32_weight_sum: intermediate_layer = intermediate_layer.astype(jnp.float32) weights = weights.astype(jnp.float32) @@ -2972,12 +3036,15 @@ def fused_moe_matmul( w0_kernel=None, w1_kernel=None, fused_kernel=None, + forced_routed_experts=None, ) -> tuple[jax.Array, None, None]: """Fused MoE via tpu_inference fused_moe_func (vllm_rpa path only). fused_moe_func handles routing, GMM, and weighted combination internally. It does not compute lb_loss or bias_updates (inference-only). """ + if forced_routed_experts is not None: + raise NotImplementedError("Forced routing via forced_routed_experts is not supported with fused_moe_matmul.") try: # pylint: disable=import-outside-toplevel # pytype: disable=import-error @@ -3071,6 +3138,7 @@ def __call__( input_ids: jax.Array | None = None, gate_inputs: jax.Array | None = None, out_sharding: NamedSharding | None = None, + forced_routed_experts: jax.Array | None = None, ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]: """Executes the routed MoE block. @@ -3134,6 +3202,7 @@ def __call__( w0_kernel=w0_kernel, w1_kernel=w1_kernel, fused_kernel=fused_kernel, + forced_routed_experts=forced_routed_experts, ) elif cfg.sparse_matmul: if quantizations.in_serve_mode(self.quant): @@ -3158,7 +3227,8 @@ def __call__( w0_bias, w1_bias, wo_bias, - input_ids, + input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) else: output, lb_loss, bias_updates = self.dense_matmul( @@ -3171,7 +3241,8 @@ def __call__( w0_bias, w1_bias, wo_bias, - input_ids, + input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) return output, lb_loss, bias_updates @@ -3261,6 +3332,7 @@ def __call__( intermediate_sharding: NamedSharding | None = None, out_sharding: NamedSharding | None = None, input_ids: jax.Array | None = None, + forced_routed_experts: jax.Array | None = None, ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]: """Executes both the routed experts and the shared expert block. @@ -3282,6 +3354,7 @@ def __call__( gate_inputs=gate_inputs, out_sharding=out_sharding, input_ids=input_ids, + forced_routed_experts=forced_routed_experts, ) shared_experts = self.shared_experts( inputs, diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ee02407502..4e324c6eb4 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1590,10 +1590,14 @@ def __call__( attention_metadata=None, deepstack_visual_embeds: None | list[jnp.ndarray] = None, multimodal_input: None | MultimodalInput = None, + forced_routed_experts: jnp.ndarray | None = None, ): cfg = self.config assert decoder_input_tokens.ndim == 2 # [batch, len] + if cfg.scan_layers and forced_routed_experts is not None: + raise NotImplementedError("Forced routing with scanned layers is not supported yet.") + policy = self.get_remat_policy() # [batch, length] -> [batch, length, emb_dim] @@ -1637,6 +1641,9 @@ def __call__( if cfg.engram_layers and decoder_input_tokens is not None: layer_kwargs["decoder_input_tokens"] = decoder_input_tokens + if forced_routed_experts is not None: + layer_kwargs["forced_routed_experts"] = forced_routed_experts + if getattr(cfg, "using_pipeline_parallelism", False): logical_partition_spec = ( self.pipeline_module.get_weight_sharding() @@ -1644,6 +1651,7 @@ def __call__( else None ) + if cfg.scan_layers: if self.is_deepseek: # Pre-pipeline: dense layers + outside-pipeline MoE layers under PP-as-DP axis rules. logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(cfg.logical_axis_rules) @@ -1768,7 +1776,6 @@ def __call__( "layer_kwargs": layer_kwargs, "decoder_input_tokens": decoder_input_tokens, } - y = self._apply_interleaved_scanned_layers( y, "dense_layers", @@ -1778,7 +1785,6 @@ def __call__( *layer_args, **common_kwargs, ) - y = self._apply_interleaved_scanned_layers( y, "moe_layers", @@ -1886,14 +1892,16 @@ def __call__( prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) dynamic_graph_init = bool(getattr(self, "disable_quant_stats_update", False)) - def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): + def pure_layer_fn(graphdef_in, state_in, y_in, kv_in, valid_kwargs=None): + if valid_kwargs is None: + valid_kwargs = layer_kwargs if cfg.parameter_memory_host_offload: state_in = jax.tree.map( lambda x: jax.device_put(x, max_utils.device_space()), state_in, ) merged_layer = nnx.merge(graphdef_in, state_in) - out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs) + out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **valid_kwargs) state_out = nnx.state(merged_layer) if dynamic_graph_init: @@ -1904,6 +1912,7 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse) + moe_lyr_idx = 0 for lyr in range(cfg.num_decoder_layers): if self.is_deepseek: if lyr < cfg.first_num_dense_layers: @@ -1943,10 +1952,38 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): if input_tokens is not None: layer_kwargs["decoder_input_tokens"] = input_tokens + current_kwargs = dict(layer_kwargs) + + is_moe = False + if cfg.decoder_block in ( + DecoderBlockType.MIXTRAL, + DecoderBlockType.QWEN3_MOE, + DecoderBlockType.QWEN3_NEXT, + DecoderBlockType.QWEN3_5, + DecoderBlockType.QWEN3_CUSTOM_MOE, + ): + is_moe = True + elif cfg.decoder_block == DecoderBlockType.DEEPSEEK: + is_moe = lyr >= cfg.first_num_dense_layers + elif cfg.decoder_block == DecoderBlockType.LLAMA4: + is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step) + + if is_moe and "forced_routed_experts" in current_kwargs and current_kwargs["forced_routed_experts"] is not None: + routed_experts = current_kwargs["forced_routed_experts"] + if routed_experts.ndim == 4: + current_kwargs["forced_routed_experts"] = routed_experts[:, :, moe_lyr_idx, :] + else: + current_kwargs["forced_routed_experts"] = routed_experts + moe_lyr_idx += 1 + else: + current_kwargs.pop("forced_routed_experts", None) + if is_moe: + moe_lyr_idx += 1 + if cfg.remat_policy != "none": - y, kv_cache, new_state, new_graphdef = checkpointed_fn(graphdef, state, y, kv_cache) + y, kv_cache, new_state, new_graphdef = checkpointed_fn(graphdef, state, y, kv_cache, current_kwargs) else: - y, kv_cache, new_state, new_graphdef = pure_layer_fn(graphdef, state, y, kv_cache) + y, kv_cache, new_state, new_graphdef = pure_layer_fn(graphdef, state, y, kv_cache, current_kwargs) if dynamic_graph_init: new_layer = nnx.merge(new_graphdef, new_state) diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index e204502cb2..06064a510b 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -34,6 +34,7 @@ import jax from jax import tree_util as jtu import qwix +import inspect M = tp.TypeVar("M", bound=Module) @@ -497,8 +498,27 @@ class ToLinen(linen.Module): # generic function to augment original nnx module (i.e for learn-to-init distillation) nnx_module_augment_fn: tp.Callable[[Module, str | None], Module] | None = None + def __post_init__(self): + super().__post_init__() + + clashing_params = {} + _call_fn = getattr(self.nnx_class, "__call__", None) + if _call_fn and callable(_call_fn): + sig = inspect.signature(_call_fn) + params = list(sig.parameters.keys()) + for name in ("previous_chunk", "page_state", "slot"): + if name in params: + clashing_params[name] = params.index(name) - 1 + object.__setattr__(self, "_clashing_params", clashing_params) + @linen.compact def __call__(self, *args, nnx_method: tp.Callable[..., Any] | str | None = None, **kwargs): + # Pop pre-partialled keyword arguments from kwargs if they are also passed positionally in args + # to avoid Python's multiple values clashing (e.g., in Linen scanned loops). + for param_name, positional_idx in getattr(self, "_clashing_params", {}).items(): + if len(args) > positional_idx: + kwargs.pop(param_name, None) + def _module_kwargs(): maybe_add_default = not self.is_initializing() module_kwargs = dict(self.kwargs) diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 0ad8978e7f..9ed51e1495 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -355,6 +355,7 @@ def __call__( kv_cache=None, attention_metadata=None, decoder_input_tokens=None, + forced_routed_experts=None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -438,6 +439,7 @@ def __call__( kv_cache=None, attention_metadata=None, decoder_input_tokens=None, + forced_routed_experts=None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -600,15 +602,20 @@ def extract_fn(x): load_balance_loss = metadata["load_balance_loss"] moe_bias_updates = metadata["moe_bias_updates"] else: - mlp_lnx, load_balance_loss, moe_bias_updates = self.mlp_op(hidden_states, deterministic) + mlp_lnx, load_balance_loss, moe_bias_updates = self.mlp_op( + hidden_states, deterministic, forced_routed_experts=forced_routed_experts + ) layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout_op(layer_output, deterministic=deterministic) return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache) - def mlp_op(self, x, deterministic, *args, **kwargs): + def mlp_op(self, x, deterministic, forced_routed_experts=None): mlp_lnx, load_balance_loss, moe_bias_updates = self.DeepSeekMoeBlock_0( - x, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding + x, + intermediate_sharding=self.mlp_intermediate_sharding, + out_sharding=self.out_sharding, + forced_routed_experts=forced_routed_experts, ) return self.with_logical_constraint(mlp_lnx), load_balance_loss, moe_bias_updates diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 14fc68cebf..e8e2895bc5 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -121,6 +121,7 @@ def __call__( original_inputs: jax.Array | None = None, intermediate_sharding: jax.sharding.NamedSharding | None = None, out_sharding: jax.sharding.NamedSharding | None = None, + forced_routed_experts: jax.Array | None = None, ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]: shared_experts = self.moe_block.shared_experts( inputs, intermediate_sharding=intermediate_sharding, out_sharding=out_sharding @@ -140,7 +141,7 @@ def __call__( # 3. Pass both to routed_moe routed_experts, load_balance_loss, moe_bias_updates = self.moe_block.routed_moe( - routed_inputs, gate_inputs=gate_inputs, out_sharding=out_sharding + routed_inputs, gate_inputs=gate_inputs, out_sharding=out_sharding, forced_routed_experts=forced_routed_experts ) routed_experts = self.post_feedforward_layernorm_2(routed_experts) @@ -321,6 +322,7 @@ def __call__( bidirectional_mask=None, kv_cache=None, attention_metadata=None, + forced_routed_experts: jnp.ndarray | None = None, ): cfg = self.config # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) @@ -365,7 +367,9 @@ def __call__( # MLP block. if getattr(self.config, "num_experts", 1) > 1: - mlp_lnx, load_balance_loss, _ = self.mlp(attn_output, original_inputs=attention_lnx) + mlp_lnx, load_balance_loss, _ = self.mlp( + attn_output, original_inputs=attention_lnx, forced_routed_experts=forced_routed_experts + ) if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) else: diff --git a/src/maxtext/models/mixtral.py b/src/maxtext/models/mixtral.py index 575bea3389..bbd0887270 100644 --- a/src/maxtext/models/mixtral.py +++ b/src/maxtext/models/mixtral.py @@ -135,6 +135,7 @@ def __call__( slot=None, kv_cache=None, attention_metadata=None, + forced_routed_experts: jnp.ndarray | None = None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -168,7 +169,7 @@ def __call__( # NOTE: the naming mismatch here is to ensure reverse compatibility with existing checkpoints. # The `name` represents the weight name in JAX/checkpoints and so the class name # is just for readability. - mlp_lnx, load_balance_loss, _ = self.MoeBlock_0(hidden_states) + mlp_lnx, load_balance_loss, _ = self.MoeBlock_0(hidden_states, forced_routed_experts=forced_routed_experts) mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) layer_output = mlp_lnx + intermediate_inputs diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index a70ad780a3..93d01a540d 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -143,6 +143,7 @@ def __call__( nnx_method=None, kv_caches: list[jax.Array] | None = None, attention_metadata: dict[str, Any] | None = None, + forced_routed_experts: jnp.ndarray | None = None, ): """Applies Transformer decoder-branch on encoded-input and target. @@ -219,6 +220,7 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, deepstack_visual_embeds=deepstack_visual_embeds, + forced_routed_experts=forced_routed_experts, ) # pytype: disable=wrong-keyword-args # If we are initializing the model AND MTP is enabled, we must create @@ -462,6 +464,7 @@ def __call__( decoder_target_mask: jax.Array | None = None, kv_caches: list[jax.Array] | None = None, attention_metadata: dict[str, Any] | None = None, + forced_routed_experts: jnp.ndarray | None = None, ): """Applies the Zero-1 FSDP wrapped Transformer model. @@ -487,6 +490,7 @@ def __call__( Returns: Logits from the Transformer model. Logits, hidden_state, kv_caches if called by vLLM. """ + if decoder_segment_ids is not None and model_mode == MODEL_MODE_AUTOREGRESSIVE: raise ValueError( f"During autoregressive decoding we assume the tokens are in the active sequence" @@ -562,6 +566,7 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, deepstack_visual_embeds=deepstack_visual_embeds, + forced_routed_experts=forced_routed_experts, ) # pytype: disable=wrong-keyword-args else: logits, hidden_state, kv_caches = self.decoder( @@ -577,6 +582,7 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, deepstack_visual_embeds=deepstack_visual_embeds, + forced_routed_experts=forced_routed_experts, mutable=mutable_collections, # pyrefly: ignore[unexpected-keyword] ) # pytype: disable=wrong-keyword-args diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 7cb710bf1e..d8eebe5d94 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1108,7 +1108,9 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn rngs=rngs, ) - def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Array | None]: + def __call__( + self, hidden_states: Array, deterministic: bool, forced_routed_experts: jnp.ndarray | None = None + ) -> tuple[Array, Array | None]: """ Applies the sparse MoE block to the input hidden states. @@ -1122,7 +1124,7 @@ def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Ar - The load balancing loss from the routed experts, if applicable during training. """ # 1. Apply the routed experts block. - routed_output, load_balance_loss, _ = self.routed_experts(hidden_states) + routed_output, load_balance_loss, _ = self.routed_experts(hidden_states, forced_routed_experts=forced_routed_experts) # 2. Apply the shared expert. shared_expert_output = self.shared_expert(hidden_states, deterministic=deterministic) @@ -1300,6 +1302,7 @@ def __call__( slot: None | int = None, kv_cache: None | dict[str, Array] = None, attention_metadata: None | dict[str, Any] = None, + forced_routed_experts: jnp.ndarray | None = None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -1342,7 +1345,9 @@ def __call__( hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) # Instantiate and call our `Qwen3NextSparseMoeBlock`. - mlp_output, load_balance_loss = self.mlp(hidden_states, deterministic=deterministic) + mlp_output, load_balance_loss = self.mlp( + hidden_states, deterministic=deterministic, forced_routed_experts=forced_routed_experts + ) # We sow the load balancing loss so it can be collected and added to the total loss # during training. @@ -1569,6 +1574,7 @@ def __call__( slot: None | int = None, kv_cache: None | jnp.ndarray = None, attention_metadata: None | dict[str, Any] = None, + forced_routed_experts: jnp.ndarray | None = None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) is_scan_carry = False @@ -1591,7 +1597,7 @@ def __call__( attention_metadata=attention_metadata, ) - mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states) + mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states, forced_routed_experts=forced_routed_experts) mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: self.moe_lb_loss = nnx.Intermediate(load_balance_loss) diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index 331df17f41..125c1ab061 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -185,6 +185,7 @@ def __call__( slot: None | int = None, kv_cache: None | dict[str, Array] = None, attention_metadata: None | dict[str, Any] = None, + forced_routed_experts: jnp.ndarray | None = None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -227,7 +228,9 @@ def __call__( hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) # Instantiate and call our `Qwen3_5SparseMoEBlock`. - mlp_output, load_balance_loss = self.mlp(hidden_states, deterministic=deterministic) + mlp_output, load_balance_loss = self.mlp( + hidden_states, deterministic=deterministic, forced_routed_experts=forced_routed_experts + ) # We sow the load balancing loss so it can be collected and added to the total loss # during training. diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index a992ca6fab..794481b257 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -153,6 +153,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr mutable=mutable_collections, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], + forced_routed_experts=data.get("forced_routed_experts", None), ) if (config.use_indexer and not config.indexer_sparse_training) and is_train: @@ -200,6 +201,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr enable_dropout=config.enable_dropout if is_train else False, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], + forced_routed_experts=data.get("forced_routed_experts", None), ) # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match # subclasses. Pop them before the generic Intermediate pop below, which would otherwise diff --git a/tests/test_trainer_router_replay.py b/tests/test_trainer_router_replay.py new file mode 100644 index 0000000000..075fc07a7e --- /dev/null +++ b/tests/test_trainer_router_replay.py @@ -0,0 +1,126 @@ +"""Integration Test: Insert and Replay Router Logits/Expert Decisions in MaxText Trainer. + +Validates: +1. Passing `forced_routed_experts` in data batch to `train.loss_fn`. +2. MoE layers executing with forced routing in trainer forward/backward step. +""" + +import os +import sys +import unittest +import jax +import jax.numpy as jnp +from jax.sharding import Mesh + +from maxtext.configs import pyconfig +from maxtext.models import models +from maxtext.utils import maxtext_utils +from maxtext.trainers.pre_train import train +from tests.utils.test_helpers import get_test_config_path + + +class TrainerRouterReplayTest(unittest.TestCase): + + def setUp(self): + os.environ["NEW_MODEL_DESIGN"] = "1" + os.environ["SKIP_JAX_PRECOMPILE"] = "1" + + def test_loss_fn_with_forced_routed_experts(self): + seq_len = 16 + batch_size = 2 + num_layers = 1 + num_experts = 4 + top_k = 2 + vocab_size = 1000 + + base_kwargs = { + "run_name": "test_trainer_router_replay", + "enable_checkpointing": False, + "override_model_config": True, + "base_num_decoder_layers": num_layers, + "num_decoder_layers": num_layers, + "model_name": "qwen3.5-35b-a3b", + "num_experts": num_experts, + "num_experts_per_tok": top_k, + "base_emb_dim": 256, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "head_dim": 256, + "partial_rotary_factor": 0.25, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "vocab_size": vocab_size, + "max_target_length": seq_len, + "max_prefill_predict_length": seq_len, + "per_device_batch_size": float(batch_size), + "scan_layers": False, + "weight_dtype": "bfloat16", + "dtype": "bfloat16", + "log_config": False, + "skip_jax_distributed_system": True, + "ici_tensor_parallelism": 1, + "ici_data_parallelism": 1, + "ici_expert_parallelism": 1, + "enable_nnx": True, + "pure_nnx": True, + "pure_nnx_decoder": True, + "sparse_matmul": True, + } + + cfg = pyconfig.initialize( + [sys.argv[0], get_test_config_path(), "attention=flash", "sparse_matmul=True"], + **base_kwargs, + ) + + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + rng = jax.random.PRNGKey(42) + + # Construct input tokens and synthetic forced routed experts + tokens = jnp.array([10, 20, 30, 40] * 4, dtype=jnp.int32)[:seq_len] + inputs = jnp.tile(jnp.expand_dims(tokens, axis=0), (batch_size, 1)) + positions = jnp.tile(jnp.expand_dims(jnp.arange(seq_len, dtype=jnp.int32), axis=0), (batch_size, 1)) + segmentation = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + targets = jnp.roll(inputs, -1, axis=-1) + + # Synthetic forced routed experts: [batch, seq_len, top_k] + forced_experts = jnp.zeros((batch_size, seq_len, top_k), dtype=jnp.int32) + forced_experts = forced_experts.at[:, :, 0].set(1) + forced_experts = forced_experts.at[:, :, 1].set(3) + + data_batch = { + "inputs": inputs, + "inputs_position": positions, + "inputs_segmentation": segmentation, + "targets": targets, + "targets_segmentation": segmentation, + "forced_routed_experts": forced_experts, + } + + model = models.transformer_as_linen(config=cfg, mesh=mesh, quant=None, model_mode="train") + init_params_rng, init_dropout_rng = jax.random.split(rng) + params = model.init( + {"params": init_params_rng, "dropout": init_dropout_rng}, + inputs, + positions, + segmentation, + enable_dropout=False, + ) + + # Execute trainer loss_fn with forced router replay data + loss, aux = train.loss_fn( + model, + cfg, + data_batch, + dropout_rng=init_dropout_rng, + params=params, + is_train=True, + ) + + self.assertIsNotNone(loss) + self.assertFalse(jnp.isnan(loss), "Loss must not be NaN") + print(f"\n[Trainer Router Replay] Computed loss with forced routing: {loss}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/forced_routing_test.py b/tests/unit/forced_routing_test.py new file mode 100644 index 0000000000..343ae885fb --- /dev/null +++ b/tests/unit/forced_routing_test.py @@ -0,0 +1,109 @@ +# 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 forced routing in moe.py.""" + +import unittest +import jax +import jax.numpy as jnp +from maxtext.layers import moe +from maxtext.common import common_types as ctypes + + +class DummyConfig: + + def __init__(self, model_name="default", decoder_block=ctypes.DecoderBlockType.DEFAULT): + self.model_name = model_name + self.decoder_block = decoder_block + self.norm_topk_prob = False + self.use_random_routing = False + self.shard_mode = ctypes.ShardMode.AUTO + + +class DummyRoutedMoE: + + def __init__(self, config): + self.config = config + self.dtype = jnp.float32 + self.num_experts_per_tok = 2 + self.num_experts = 3 + + def _maybe_shard_with_logical(self, x, spec): + return x + + +class ForcedRoutingTest(unittest.TestCase): + + def test_basic_override(self): + config = DummyConfig() + model = DummyRoutedMoE(config) + + gate_logits = jnp.array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]) # (1, 2, 3) + pre_bias_logits = gate_logits # Not DeepSeek + forced_routed_experts = jnp.array([[[2, 1], [0, 2]]]) # (1, 2, 2) + + top_k_weights, top_k_indices = moe.RoutedMoE.get_topk( + model, gate_logits, pre_bias_logits, forced_routed_experts=forced_routed_experts + ) + + # Check that indices are overridden + self.assertTrue((top_k_indices == forced_routed_experts).all()) + # Check that weights are extracted correctly and softmaxed + # For token 0: indices 2, 1 -> logits 3.0, 2.0 -> softmax([3.0, 2.0]) + # For token 1: indices 0, 2 -> logits 4.0, 6.0 -> softmax([4.0, 6.0]) + expected_weights = jax.nn.softmax(jnp.array([[[3.0, 2.0], [4.0, 6.0]]]).astype(jnp.float32), axis=-1) + self.assertTrue(jax.numpy.allclose(top_k_weights, expected_weights, rtol=1e-5, atol=1e-5)) + + def test_gemma4_softmax(self): + config = DummyConfig(decoder_block=ctypes.DecoderBlockType.GEMMA4) + model = DummyRoutedMoE(config) + + gate_logits = jnp.array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]) # (1, 2, 3) + pre_bias_logits = gate_logits + forced_routed_experts = jnp.array([[[2, 1], [0, 2]]]) # (1, 2, 2) + + top_k_weights, top_k_indices = moe.RoutedMoE.get_topk( + model, gate_logits, pre_bias_logits, forced_routed_experts=forced_routed_experts + ) + + # Check that indices are overridden + self.assertTrue((top_k_indices == forced_routed_experts).all()) + + # For Gemma 4, it applies softmax to gate_logits first! + + expected_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) + expected_weights = jnp.take_along_axis(expected_probs, forced_routed_experts, axis=-1) + + self.assertTrue(jax.numpy.allclose(top_k_weights, expected_weights, rtol=1e-5, atol=1e-5)) + + def test_reshape_and_update_weights(self): + config = DummyConfig() + model = DummyRoutedMoE(config) + + weights = jnp.array([[[0.1, 0.2], [0.3, 0.4]]]) # (1, 2, 2) + indices = jnp.array([[[2, -1], [-1, 1]]]) # (1, 2, 2) + + update_weights = moe.RoutedMoE.reshape_and_update_weights(model, weights, indices, safe_updates=True) + + # Expected shape: (1, 2, 3) where 3 is num_experts! + # For token 0: index 2 -> 0.1. Index -1 -> mapped to 0 but weight 0.0! + # So for expert 0: 0.0. Expert 1: 0.0. Expert 2: 0.1. + # For token 1: index -1 -> mapped to 0 but weight 0.0! Index 1 -> 0.4. + # So for expert 0: 0.0. Expert 1: 0.4. Expert 2: 0.0. + expected_update_weights = jnp.array([[[0.0, 0.0, 0.1], [0.0, 0.4, 0.0]]]) + + self.assertTrue(jax.numpy.allclose(update_weights, expected_update_weights, rtol=1e-5, atol=1e-5)) + + +if __name__ == "__main__": + unittest.main()