Measure routing mismatch for Qwen and DeepSeek models

Author: xuefgu

src/maxtext/integration/tunix/tunix_adapter.py

@@ -70,6 +70,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.
@@ -80,6 +81,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
 

src/maxtext/layers/decoders.py

@@ -793,9 +793,12 @@ def __call__(
       kv_caches: list[jax.Array] | None = None,
       attention_metadata=None,
       deepstack_visual_embeds: None | list[jnp.ndarray] = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     cfg = self.config
     mesh = self.mesh
+    if cfg.use_batch_split_schedule and forced_routed_experts is not None:
+      raise NotImplementedError("Forced routing mismatch measurement is not supported with batch split schedule.")
     assert decoder_input_tokens.ndim == 2  # [batch, len]
 
     # [batch, length] -> [batch, length, emb_dim]
@@ -882,24 +885,44 @@ def __call__(
           }
           dense_layer = RemattedBlockLayers[0]
           moe_layer = RemattedBlockLayers[1]
-          if cfg.engram_layers:
-            original_dense_call = dense_layer.__call__
-            original_moe_call = moe_layer.__call__
-            dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
-            moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)
 
+          dense_broadcast_args = list(broadcast_args)
+          dense_in_axes_tuple = [nn.broadcast] * len(broadcast_args)
+          input_tokens = decoder_input_tokens if cfg.engram_layers else None
+          dense_broadcast_args.extend([previous_chunk, slot, None, attention_metadata, input_tokens])
+          dense_in_axes_tuple.extend([nn.broadcast] * 5)
+
+          moe_broadcast_args = list(broadcast_args)
+          moe_in_axes_tuple = [nn.broadcast] * len(broadcast_args)
+          moe_broadcast_args.extend([previous_chunk, slot, None, attention_metadata, input_tokens])
+          moe_in_axes_tuple.extend([nn.broadcast] * 5)
+
+          num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
+
+          forced_routed_experts_t = None
+          if forced_routed_experts is not None:
+            forced_routed_experts_t = jnp.transpose(forced_routed_experts, (2, 0, 1, 3))
+            moe_broadcast_args.append(forced_routed_experts_t)
+            moe_in_axes_tuple.append(0)
+          else:
+            moe_broadcast_args.append(None)
+            moe_in_axes_tuple.append(nn.broadcast)
+
+          if cfg.engram_layers:
             common_kwargs = {
                 "dense_layer": dense_layer,
                 "moe_layer": moe_layer,
-                "original_dense_call": original_dense_call,
-                "original_moe_call": original_moe_call,
                 "layer_call_kwargs": layer_call_kwargs,
                 "decoder_segment_ids": decoder_segment_ids,
                 "decoder_positions": decoder_positions,
                 "deterministic": deterministic,
                 "model_mode": model_mode,
                 "decoder_input_tokens": decoder_input_tokens,
-                "broadcast_args": broadcast_args,
+                "dense_broadcast_args": dense_broadcast_args,
+                "dense_in_axes": dense_in_axes_tuple,
+                "moe_broadcast_args": moe_broadcast_args,
+                "moe_in_axes": moe_in_axes_tuple,
+                "forced_routed_experts": forced_routed_experts_t,
             }
 
             # Apply Dense Layers
@@ -922,23 +945,22 @@ def __call__(
                 **common_kwargs,
             )
           else:
-            dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
             y, _ = self.scan_decoder_layers(
                 cfg,
                 dense_layer,
                 cfg.first_num_dense_layers,
                 "dense_layers",
                 mesh,
-                in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
+                in_axes_tuple=tuple(dense_in_axes_tuple),
                 model_mode=model_mode,
-            )(y, *broadcast_args)
-            moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)
-            num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
+            )(y, *dense_broadcast_args)
 
             # If batch-split schedule is used and initialization is complete,
             # as detected by immutable params, use deepseek_batchsplit custom
             # scan with initialized parameters.
             if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"):
+              if forced_routed_experts is not None:
+                raise NotImplementedError("Forced routing mismatch measurement is not supported with batch split schedule.")
               # old version of batch-split that fully uses qwix quantization.
               if cfg.use_qwix_quantization and not cfg.use_manual_quantization:
                 y = deepseek_batchsplit_fp8.scan_batch_split_layers(
@@ -971,9 +993,9 @@ def __call__(
                   num_moe_layers,
                   "moe_layers",
                   mesh,
-                  in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
+                  in_axes_tuple=tuple(moe_in_axes_tuple),
                   model_mode=model_mode,
-              )(y, *broadcast_args)
+              )(y, *moe_broadcast_args)
         elif cfg.decoder_block == DecoderBlockType.GEMMA3:
           bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None
           y = self._apply_gemma3_scanned_blocks(
@@ -1012,19 +1034,35 @@ def __call__(
           current_broadcast_args = list(broadcast_args)
           current_in_axes_tuple = list(in_axes_tuple)
 
+          if forced_routed_experts is not None:
+            # Transpose [B, L, N, E] -> [N, B, L, E] for scan
+            forced_routed_experts_t = jnp.transpose(forced_routed_experts, (2, 0, 1, 3))
+            cycle_interval = cfg.inhomogeneous_layer_cycle_interval
+            reshaped_forced = jnp.reshape(
+                forced_routed_experts_t,
+                (scan_length, cycle_interval) + forced_routed_experts_t.shape[1:]
+            )
+          else:
+            reshaped_forced = None
+
           if kv_caches is not None:
             # Stack kv_caches for scan: [num_layers, ...]
             stacked_kv_cache = jnp.stack(kv_caches, axis=0)
 
             # We pass (y, stacked_kv_cache, 0) as the carry
             carry = (y, stacked_kv_cache, 0)
 
-            # We don't pass kv_cache as a scanned argument anymore
-
             # Pass None for previous_chunk, slot, kv_cache to align with __call__ signature
             current_broadcast_args.extend([None, None, None, attention_metadata])
             current_in_axes_tuple.extend([nn.broadcast] * 4)
 
+            if reshaped_forced is not None:
+              current_broadcast_args.append(reshaped_forced)
+              current_in_axes_tuple.append(0)
+            else:
+              current_broadcast_args.append(None)
+              current_in_axes_tuple.append(nn.broadcast)
+
             max_logging.info(f"DEBUG: len(current_broadcast_args)={len(current_broadcast_args)}")
             max_logging.info(f"DEBUG: current_broadcast_args={[type(a) for a in current_broadcast_args]}")
 
@@ -1046,8 +1084,15 @@ def __call__(
               kv_caches[i] = returned_kv_cache[i]
           else:
             # Fallback to old behavior if kv_caches is None (not vLLM RPA)
-            current_broadcast_args.append(None)
-            current_in_axes_tuple.append(nn.broadcast)
+            current_broadcast_args.extend([None, None, None, None])
+            current_in_axes_tuple.extend([nn.broadcast] * 4)
+
+            if reshaped_forced is not None:
+              current_broadcast_args.append(reshaped_forced)
+              current_in_axes_tuple.append(0)
+            else:
+              current_broadcast_args.append(None)
+              current_in_axes_tuple.append(nn.broadcast)
 
             y, _ = self.scan_decoder_layers(
                 cfg,
@@ -1076,6 +1121,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
+              extra_kwargs = {}
+              if layer_prefix == "moe_layers" and forced_routed_experts is not None:
+                extra_kwargs["forced_routed_experts"] = forced_routed_experts[:, :, index, :]
+
               y, kv_cache = layer(
                   config=cfg,
                   mesh=mesh,
@@ -1094,6 +1143,7 @@ def __call__(
                   kv_cache=kv_cache,
                   attention_metadata=attention_metadata,
                   decoder_input_tokens=input_tokens,
+                  **extra_kwargs,
               )
               if kv_caches is not None and kv_cache is not None:
                 kv_caches[index] = kv_cache
@@ -1113,6 +1163,7 @@ def __call__(
               slot=slot,
           )
         else:
+          moe_lyr_idx = 0
           for lyr in range(cfg.num_decoder_layers):
             RemattedBlockLayer = RemattedBlockLayers[0]
             layer_kwargs = {}
@@ -1149,6 +1200,30 @@ def __call__(
             layer = RemattedBlockLayer(
                 config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **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.LLAMA4:
+              is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step)
+
+            current_forced_routed_experts = None
+            if is_moe and forced_routed_experts is not None:
+              current_forced_routed_experts = forced_routed_experts[:, :, moe_lyr_idx, :]
+              moe_lyr_idx += 1
+            elif is_moe:
+              moe_lyr_idx += 1
+
+            extra_kwargs = {}
+            if is_moe and current_forced_routed_experts is not None:
+              extra_kwargs["forced_routed_experts"] = current_forced_routed_experts
+
             y, returned_cache = layer(
                 y,
                 decoder_segment_ids,
@@ -1159,6 +1234,7 @@ def __call__(
                 slot=slot,
                 kv_cache=kv_cache,
                 attention_metadata=attention_metadata,
+                **extra_kwargs,
                 **layer_call_kwargs,
             )
             if kv_caches is not None and returned_cache is not None:
@@ -1466,10 +1542,12 @@ def _apply_single_engram_layer(self, y, current_idx, layer_type, **kwargs):
     """Applies a single, unscanned Engram layer."""
     layer = kwargs["dense_layer"] if layer_type == "dense" else kwargs["moe_layer"]
     layer_prefix = "dense_layers" if layer_type == "dense" else "moe_layers"
-    original_call = kwargs["original_dense_call"] if layer_type == "dense" else kwargs["original_moe_call"]
-    layer_call_kwargs = kwargs["layer_call_kwargs"]
 
-    layer.__call__ = original_call
+    extra_kwargs = {}
+    if layer_type == "moe" and kwargs.get("forced_routed_experts") is not None:
+      moe_idx = current_idx - self.config.first_num_dense_layers
+      extra_kwargs["forced_routed_experts"] = kwargs["forced_routed_experts"][moe_idx]
+
     y, _ = layer(
         config=self.config,
         mesh=self.mesh,
@@ -1483,27 +1561,35 @@ def _apply_single_engram_layer(self, y, current_idx, layer_type, **kwargs):
         kwargs["decoder_positions"],
         kwargs["deterministic"],
         kwargs["model_mode"],
+        previous_chunk=kwargs["layer_call_kwargs"]["previous_chunk"],
+        slot=kwargs["layer_call_kwargs"]["slot"],
         decoder_input_tokens=kwargs["decoder_input_tokens"],
-        **layer_call_kwargs,
+        **extra_kwargs,
     )
-    layer.__call__ = functools.partial(original_call, **layer_call_kwargs)
     return y
 
   def _apply_scanned_chunk(self, y, current_idx, next_boundary, layer_type, **kwargs):
     """Applies a contiguous chunk of layers using the scan operation."""
     layer = kwargs["dense_layer"] if layer_type == "dense" else kwargs["moe_layer"]
     layer_prefix = "dense_layers" if layer_type == "dense" else "moe_layers"
-    broadcast_args = kwargs["broadcast_args"]
+    broadcast_args = kwargs["dense_broadcast_args"] if layer_type == "dense" else list(kwargs["moe_broadcast_args"])
+    in_axes_tuple = kwargs["dense_in_axes"] if layer_type == "dense" else list(kwargs["moe_in_axes"])
     scan_length = next_boundary - current_idx
 
     if scan_length > 0:
+      if layer_type == "moe" and broadcast_args[-1] is not None:
+        moe_start = current_idx - self.config.first_num_dense_layers
+        moe_end = next_boundary - self.config.first_num_dense_layers
+        sliced_forced = broadcast_args[-1][moe_start:moe_end]
+        broadcast_args[-1] = sliced_forced
+
       y, _ = self.scan_decoder_layers(
           self.config,
           layer,
           scan_length,
           f"{layer_prefix}_{current_idx}_{next_boundary - 1}",
           self.mesh,
-          in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
+          in_axes_tuple=tuple(in_axes_tuple),
           model_mode=kwargs["model_mode"],
       )(y, *broadcast_args)
     return y