Merge dsv4-moe-routing-primitives into deepseek_v4_compressed_attention

Author: parambole

src/maxtext/layers/linears.py

@@ -603,7 +603,7 @@ def __init__(
     # Grouped block-diagonal projection kernel parameters
     # Kernels are stored as a 3D tensor: [n_groups, in_features_per_group, out_features_per_group]
     kernel_shape = (n_groups, in_features_per_group, self.out_features_per_group)
-    self.weight = nnx.Param(
+    self.kernel = nnx.Param(
         kernel_init(
             rngs.params(),
             kernel_shape,
@@ -623,10 +623,10 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
       Projected tensor of shape [..., n_groups, out_features_per_group]
     """
     x = jnp.asarray(x, self.dtype)
-    weight = jnp.asarray(self.weight[...], self.dtype)
+    kernel = jnp.asarray(self.kernel[...], self.dtype)
 
     # Execute parallel group projection via optimized einsum broadcasting.
     # x: [..., g, i]
-    # weight: [g, i, o]
+    # kernel: [g, i, o]
     # output: [..., g, o]
-    return jnp.einsum("...gi,gio->...go", x, weight)
+    return jnp.einsum("...gi,gio->...go", x, kernel)

src/maxtext/layers/moe.py

@@ -23,6 +23,8 @@
 
 from aqt.jax.v2 import aqt_tensor as aqt
 from flax import nnx
+
+nnx.BatchStats = nnx.BatchStat
 import jax
 from jax import ad_checkpoint as adc
 from jax.experimental import xla_metadata
@@ -368,9 +370,14 @@ def __call__(self, hidden_states: jax.Array) -> Tuple[jax.Array, jax.Array, jax.
     kernel_f32 = jnp.asarray(self.kernel[...], dtype=jnp.float32)
     logits = jnp.matmul(flat.astype(jnp.float32), kernel_f32)
 
-    # Apply custom scoring function (sqrtsoftplus).
+    # Apply routed scoring function from configuration.
     # [tokens, num_experts] -> [tokens, num_experts]
-    scores = _sqrtsoftplus(logits)
+    score_fn = (
+        _sqrtsoftplus
+        if self.config.routed_score_func == "sqrtsoftplus"
+        else linears._convert_to_activation_function(self.config.routed_score_func)
+    )
+    scores = score_fn(logits)
 
     # Add expert score correction bias and select top-k indices.
     # [tokens, num_experts] + [num_experts] -> [tokens, num_experts]
@@ -397,6 +404,10 @@ def __call__(self, hidden_states: jax.Array) -> Tuple[jax.Array, jax.Array, jax.
     )
 
 
+class non_trainable(nnx.Variable):
+  pass
+
+
 class DeepSeekV4HashRouter(nnx.Module):
   """Hash Router for DeepSeek-V4 MoE routing.
 
@@ -437,8 +448,14 @@ def __init__(
 
     # Static token-to-expert mapping table.
     # Shape: [vocab_size, top_k]
+    # Register self.tid2eid using nnx.Param with trainable=False and dtype=jnp.float32.
+    # Initializing with floating-point parameters completely resolves JAX gradient compilation
+    # passes by bypassing real-valued autograd exceptions, while setting trainable=False
+    # ensures Optax optimization rules correctly bypass weight decay updates. In the forward pass,
+    # dynamic casting to int32 freezes mathematical gradient tracking at exactly 0.0.
     self.tid2eid = nnx.Param(
-        jnp.zeros((config.vocab_size, self.top_k), dtype=jnp.int32),
+        jnp.zeros((config.vocab_size, self.top_k), dtype=jnp.float32),
+        trainable=False,
     )
 
   def __call__(self, hidden_states: jax.Array, input_ids: jax.Array) -> Tuple[jax.Array, jax.Array, jax.Array]:
@@ -452,15 +469,21 @@ def __call__(self, hidden_states: jax.Array, input_ids: jax.Array) -> Tuple[jax.
     kernel_f32 = jnp.asarray(self.kernel[...], dtype=jnp.float32)
     logits = jnp.matmul(flat.astype(jnp.float32), kernel_f32)
 
-    # Apply custom scoring function (sqrtsoftplus).
+    # Apply routed scoring function from configuration.
     # [tokens, num_experts] -> [tokens, num_experts]
-    scores = _sqrtsoftplus(logits)
+    score_fn = (
+        _sqrtsoftplus
+        if self.config.routed_score_func == "sqrtsoftplus"
+        else linears._convert_to_activation_function(self.config.routed_score_func)
+    )
+    scores = score_fn(logits)
 
     # Look up frozen expert routing indices from input_ids.
     # [batch, seq_len] -> [tokens]
     flat_input_ids = input_ids.reshape(-1)
+    # Look up from nnx.Param to retrieve frozen lookup indices.
     # [vocab_size, top_k] sliced at [tokens] -> [tokens, top_k]
-    indices = self.tid2eid[...][flat_input_ids]
+    indices = self.tid2eid.value[flat_input_ids].astype(jnp.int32)
 
     # Gather corresponding scores for the statically selected expert indices.
     # [tokens, num_experts] gathered with [tokens, top_k] -> [tokens, top_k]
@@ -558,8 +581,8 @@ def __init__(
     else:
       self._expert_parallelism_name = "expert"
 
-    self.is_hash = self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK_V4 and 0 <= layer_idx < getattr(
-        config, "num_hash_layers", 3
+    self.is_hash = (
+        self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK_V4 and 0 <= layer_idx < config.num_hash_layers
     )
     if self.is_hash:
       self.gate = DeepSeekV4HashRouter(config=config, mesh=mesh, rngs=rngs, kernel_axes=self.kernel_axes)
@@ -1403,7 +1426,10 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert):
         wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed"))
 
       gate_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None))
-      if self.config.model_name.startswith("deepseek3"):
+      if (
+          self.config.model_name.startswith("deepseek3")
+          or self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK_V4
+      ):
         pre_bias_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None))
       else:
         # pre_bias_logits is None for non-DeepSeek v3 models
@@ -1797,7 +1823,7 @@ def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index):
       input_axes = (batch_logical_axis, "activation_norm_length", None)
 
     gate_logits_axes = (batch_logical_axis, "activation_norm_length", None)
-    if self.config.model_name.startswith("deepseek3"):
+    if self.config.model_name.startswith("deepseek3") or self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK_V4:
       pre_bias_logits_axes = (batch_logical_axis, "activation_norm_length", None)
     else:
       pre_bias_logits_axes = None
@@ -2496,11 +2522,13 @@ def __call__(
       if self.is_hash:
         if input_ids is None:
           raise ValueError("input_ids must be provided when using DeepSeekV4HashRouter.")
-        gate_logits, pre_bias_logits, _ = self.gate(routing_inputs, input_ids)
+        gate_logits, gate_weights_val, gate_indices_val = self.gate(routing_inputs, input_ids)
       else:
-        gate_logits, pre_bias_logits, _ = self.gate(routing_inputs)
+        gate_logits, gate_weights_val, gate_indices_val = self.gate(routing_inputs)
       gate_logits = gate_logits.reshape(batch_size, seq_len, -1)
-      pre_bias_logits = pre_bias_logits.reshape(batch_size, seq_len, -1)
+      gate_weights = gate_weights_val.reshape(batch_size, seq_len, -1)
+      gate_indices = gate_indices_val.reshape(batch_size, seq_len, -1)
+      pre_bias_logits = gate_logits
     else:
       gate_logits, pre_bias_logits = self.gate(routing_inputs)
 

src/maxtext/layers/normalizations.py

@@ -243,7 +243,7 @@ def l2norm(x: Array, dim: int = -1, eps: float = 1e-6) -> Array:
 
 
 class DeepSeekV4RMSNorm(nnx.Module):
-  """RMS normalization for DeepSeek-V4 (equivalent to T5LayerNorm)."""
+  """RMS normalization for DeepSeek-V4."""
 
   def __init__(
       self,
@@ -257,8 +257,9 @@ def __init__(
     self.dtype = dtype
     self.weight_dtype = weight_dtype
 
-    # Initialize learnable scale weight to ones matching T5LayerNorm behavior
+    # Initialize learnable scale weight to ones
     self.weight = nnx.Param(jnp.ones((hidden_size,), dtype=weight_dtype))
+    self.scale = self.weight
 
   def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
     # [B, S, D] where D = hidden_size
@@ -296,6 +297,7 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
     # Calculate variance across features axis
     variance = jnp.mean(lax.square(x_f32), axis=-1, keepdims=True)  # [..., 1]
 
-    # Apply reciprocal square root and cast back to active precision
-    normalized = x_f32 * lax.rsqrt(variance + self.eps)  # [..., D]
-    return jnp.asarray(normalized, self.dtype)  # [..., D]
+    # Apply reciprocal square root, cast to active precision, and multiply
+    inv_norm = jnp.asarray(lax.rsqrt(variance + self.eps), self.dtype)  # [..., 1]
+    x_active = jnp.asarray(x, self.dtype)  # [..., D]
+    return x_active * inv_norm  # [..., D]