Emit the LM head weight gradient in kernel order under explicit sharding

docs/guides/optimization.md

@@ -49,6 +49,13 @@ Optimize with Pallas kernels for fine-grained control.
 
 Guide to setting up benchmarks, performing performance tuning, and analyzing metrics.
 ```
+
+```{grid-item-card} 🔀 Shard Mode: explicit vs auto
+:link: optimization/shard_mode_performance
+:link-type: doc
+
+Measured HLO and xprof comparison of `shard_mode: explicit` against `shard_mode: auto` on the onboarded models.
+```
 ````
 
 ```{toctree}
@@ -61,4 +68,5 @@ optimization/sharding.md
 optimization/custom_mesh_and_rule.md
 optimization/pallas_kernels_performance.md
 optimization/benchmark_and_performance.md
+optimization/shard_mode_performance.md
 ```

src/maxtext/configs/base.yml

@@ -200,6 +200,15 @@ logits_via_embedding: false
 normalize_embedding_logits: true  # whether to normalize pre-softmax logits if logits_via_embedding is true
 logits_dot_in_fp32: false  # whether to use fp32 in logits_dense or shared_embedding dot product for stability
 cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher precision is generally beneficial, but it can vary slightly.
+# Trace the untied LM head's dot inside a jax.sharding.auto_axes region so that XLA folds the transpose back into
+# the gradient dot and the weight gradient comes out in the kernel's stored orientation. Without the region that
+# fold never happens under shard_mode: explicit, because the Sharding custom-call on every dot output is a barrier
+# algsimp will not cross, so the gradient is left sharded on its minor-most dimension and loses the all-reduce ->
+# reduce-scatter fusion (see docs/guides/optimization/shard_mode_performance.md). The stored kernel, its
+# initialization and the arithmetic are untouched, so gradients are bit-identical to the default rule and no
+# checkpoint conversion is needed. No effect under shard_mode: auto, where XLA folds the transpose itself, and none
+# on a tied head, which has no kernel of its own -- so None means on exactly where it can do something.
+lm_head_weight_grad_in_kernel_order: None  # on for an untied model under shard_mode: explicit if None
 float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product
 float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax
 mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively.
@@ -710,6 +719,16 @@ shard_optimizer_over_data: false
 # when dense MLP weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
 dense_fsdp_use_two_stage_all_gather: false
 
+# The in-loop counterpart of lm_head_weight_grad_in_kernel_order: trace the attention q/k/v/out and MLP wi/wo dots
+# inside a jax.sharding.auto_axes region so that XLA folds the transpose back into the gradient dot and each weight
+# gradient comes out in its kernel's stored orientation. Same barrier as the LM head flag, same default rule -- None
+# means on wherever it can act, i.e. under shard_mode: explicit. Without it explicit is 0.33%-1.48% slower than auto
+# at every measured layer count that is not a multiple of 8; with it explicit lands within 0.2% of auto everywhere,
+# and the two flags are additive (see docs/guides/optimization/shard_mode_performance.md section 4.8). Stored kernels,
+# their initialization and the arithmetic are untouched, so gradients are bit-identical to the default rule and no
+# checkpoint conversion is needed. No effect under shard_mode: auto or on quantized layers.
+dense_weight_grad_in_kernel_order: None  # on under shard_mode: explicit if None
+
 # Unless explicitly specified, the number of TPU slices is automatically determined. It should only be set for
 # disaggregated reinforcement learning workloads using multiple slices. For ahead of time compilation,
 # you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1.

src/maxtext/configs/types.py

@@ -604,6 +604,16 @@ class LogitsAndLoss(BaseModel):
   )
   logits_dot_in_fp32: bool = Field(False, description="Use fp32 for the logits dot product for stability.")
   cast_logits_to_fp32: bool = Field(True, description="Whether to cast the final logits to fp32.")
+  lm_head_weight_grad_in_kernel_order: bool | None = Field(
+      None,
+      description=(
+          "Trace the untied LM head's dot inside a jax.sharding.auto_axes region so that XLA folds the "
+          "transpose back into the gradient dot and the weight gradient comes out in the kernel's stored "
+          "axis order. Recovers the weight-gradient reduce-scatter that shard_mode=explicit otherwise loses, "
+          "leaving the stored kernel, its initialization and the arithmetic untouched. No effect under "
+          "shard_mode=auto. None means on for an untied model under shard_mode=explicit and off everywhere else."
+      ),
+  )
   final_logits_soft_cap: None | NonNegativeFloat = Field(
       None,
       description="Soft-cap value for the final logits. None or 0.0 means no cap.",
@@ -1258,6 +1268,20 @@ class LayoutAndSharding(BaseModel):
       False,
       description="Use two separate All-Gather calls for dense MLP weights sharded on both FSDP and FSDP-transpose.",
   )
+  dense_weight_grad_in_kernel_order: bool | None = Field(
+      None,
+      description=(
+          "Trace the per-layer attention and MLP projection dots inside a jax.sharding.auto_axes region so that "
+          "XLA folds the transpose back into the gradient dot and each weight gradient comes out in its kernel's "
+          "stored axis order. This is the in-loop counterpart of lm_head_weight_grad_in_kernel_order and, like it, "
+          "None means on wherever it can act: shard_mode=explicit, where the Sharding custom-call on every dot "
+          "output blocks that fold. Without it explicit is 0.33%-1.48% slower than auto at every measured layer "
+          "count that is not a multiple of 8; with it explicit lands within 0.2% of auto everywhere "
+          "(docs/guides/optimization/shard_mode_performance.md section 4.8). Stored kernels, their initialization "
+          "and the arithmetic are untouched, so gradients are bit-identical to the default rule. No effect under "
+          "shard_mode=auto, where XLA folds the transpose itself, or on quantized layers."
+      ),
+  )
   internal_compile: bool = Field(
       False,
       description="Use internal_compile to bypass open-source topology mappings.",
@@ -3269,6 +3293,39 @@ def validate_shard_embed_moe_on_fsdp(self) -> "MaxTextConfig":
       )
     return self
 
+  @model_validator(mode="after")
+  def resolve_lm_head_weight_grad_in_kernel_order(self) -> "MaxTextConfig":
+    """Resolve the tri-state flag, and reject it where it cannot be honored.
+
+    The transpose it removes only exists under explicit sharding, and on an untied
+    model removing it has been a win or a wash on every configuration measured
+    (docs/guides/optimization/shard_mode_performance.md section 5) -- on qwen3-8b it
+    is the difference between +3.25% and -0.41% against `auto`. So the default is
+    "on wherever it can do anything", and writing the flag out is only needed to
+    reproduce a measurement.
+    """
+    if self.lm_head_weight_grad_in_kernel_order is None:
+      self.lm_head_weight_grad_in_kernel_order = self.shard_mode == ShardMode.EXPLICIT and not self.logits_via_embedding
+    elif self.lm_head_weight_grad_in_kernel_order and self.logits_via_embedding:
+      raise ValueError(
+          "lm_head_weight_grad_in_kernel_order only applies to the untied LM head, but logits_via_embedding is True."
+      )
+    return self
+
+  @model_validator(mode="after")
+  def resolve_dense_weight_grad_in_kernel_order(self) -> "MaxTextConfig":
+    """Resolve the in-loop counterpart of the LM-head flag.
+
+    Same barrier, same shape of fix, same default rule: on wherever it can act.
+    The two are independent and additive -- on llama2 at 18 layers, explicit
+    costs +1.098% against auto with neither flag, +0.621% with this one alone,
+    +0.429% with the LM-head flag alone and +0.005% with both
+    (docs/guides/optimization/shard_mode_performance.md section 4.8).
+    """
+    if self.dense_weight_grad_in_kernel_order is None:
+      self.dense_weight_grad_in_kernel_order = self.shard_mode == ShardMode.EXPLICIT
+    return self
+
   @model_validator(mode="after")
   def set_derived_and_validate_values(self) -> "MaxTextConfig":
     """

src/maxtext/layers/attention_mla.py

@@ -138,6 +138,7 @@ def __init__(
         quant=self.quant,
         matmul_precision=self.config.matmul_precision,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         rngs=self.rngs,
     )
 
@@ -153,6 +154,7 @@ def __init__(
         quant=self.quant,
         matmul_precision=self.config.matmul_precision,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         rngs=self.rngs,
     )
 
@@ -172,6 +174,7 @@ def __init__(
         quant=None,
         matmul_precision=self.config.matmul_precision,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         rngs=self.rngs,
     )
 
@@ -808,6 +811,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No
           quant=self.quant,
           matmul_precision=self.config.matmul_precision,
           shard_mode=self.config.shard_mode,
+          weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
           rngs=self.rngs,
       )
     else:
@@ -823,6 +827,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No
           quant=self.quant,
           matmul_precision=self.config.matmul_precision,
           shard_mode=self.config.shard_mode,
+          weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
           rngs=self.rngs,
       )
       self.q_norm = RMSNorm(
@@ -844,6 +849,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No
           quant=self.quant,
           matmul_precision=self.config.matmul_precision,
           shard_mode=self.config.shard_mode,
+          weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
           rngs=self.rngs,
       )
 
@@ -859,6 +865,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No
         quant=self.quant,
         matmul_precision=self.config.matmul_precision,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         rngs=self.rngs,
     )
     self.kv_norm = RMSNorm(
@@ -883,6 +890,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No
         quant=self.quant,
         matmul_precision=self.config.matmul_precision,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         rngs=self.rngs,
     )
 

src/maxtext/layers/attentions.py

@@ -697,6 +697,7 @@ def query_init(*args):
         matmul_precision=self.config.matmul_precision,
         use_bias=self.use_bias_in_projections,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         rngs=self.rngs,
     )
 
@@ -733,6 +734,7 @@ def init_kv_w(self, inputs_kv_shape: Tuple) -> nnx.Module:
         weight_dtype=self.weight_dtype,
         quant=self.quant,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         matmul_precision=self.config.matmul_precision,
         use_bias=self.use_bias_in_projections,
         rngs=self.rngs,
@@ -774,6 +776,7 @@ def init_qkv_w(self, inputs_shape: Tuple) -> nnx.Module:
         weight_dtype=self.weight_dtype,
         quant=self.quant,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         matmul_precision=self.config.matmul_precision,
         use_bias=self.use_bias_in_projections,
         rngs=self.rngs,
@@ -829,6 +832,7 @@ def init_out_w(self, output_dim: int) -> nnx.Module:
         weight_dtype=self.weight_dtype,
         quant=self.quant,
         shard_mode=self.config.shard_mode,
+        weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order,
         matmul_precision=self.config.matmul_precision,
         use_bias=False if self.is_qwen2 else self.use_bias_in_projections,
         rngs=self.rngs,

src/maxtext/layers/decoders.py

@@ -838,6 +838,7 @@ def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, determi
           name="logits_dense",
           matmul_precision=self.config.matmul_precision,
           parameter_memory_host_offload=cfg.parameter_memory_host_offload,
+          weight_grad_in_kernel_order=cfg.lm_head_weight_grad_in_kernel_order,
       )(
           y,
           out_sharding=out_sharding,

src/maxtext/layers/embeddings.py

@@ -35,6 +35,20 @@
 _MAX_WAVELENGTH = 10_000
 
 
+def _cis(theta: Array) -> Array:
+  """`exp(1j * theta)` for real `theta`, written as `cos(theta) + 1j*sin(theta)`.
+
+  Bit-identical to `jnp.exp(1j * theta)`, but it does not go through XLA's
+  overflow-safe complex `exponential` expansion, which evaluates
+  `exp(real(1j*theta)) == exp(0)` over the whole tensor. Under
+  `shard_mode: explicit` it evaluates it *twice*: the `Sharding` custom-call on
+  the broadcast `1j` keeps the simplifier from commuting the constant to the
+  right, so the reassociation that lets CSE merge the two `exponential`s never
+  fires.
+  """
+  return jax.lax.complex(jnp.cos(theta), jnp.sin(theta))
+
+
 def _maybe_move_embedding_to_device(embedding_table: Array, config: Config) -> Array:
   """Moves embedding table to device if parameter offloading is enabled."""
   if config.parameter_memory_host_offload:
@@ -232,9 +246,18 @@ def attend_on_embedding(
   if out_sharding is not None:
     out_sharding = truncate_out_sharding(out_sharding, query.ndim)
   embedding_table = _maybe_move_embedding_to_device(embedding_table, config)
-  return jnp.dot(
+  # Contract over the table's feature axis instead of materializing `table.T`.
+  # Under `shard_mode: explicit` the transposed table is a distinct typed value,
+  # so XLA cannot fold it into the dot's dimension numbers: the (large) table
+  # shard is cast to bf16 once for the input lookup in `Embed.__call__` and a
+  # second time here, and the tied head's weight gradient comes out flipped.
+  # Expressing the transpose as dimension numbers keeps both consumers on one
+  # cast and leaves the gradient in the table's own axis order. Under `auto`
+  # this is a no-op -- XLA already folded the `.T` away.
+  return jnp.einsum(
+      "...e,ve->...v",
       query,
-      jnp.asarray(embedding_table, jnp.bfloat16).T,
+      jnp.asarray(embedding_table, jnp.bfloat16),
       preferred_element_type=attend_dtype,
       out_sharding=out_sharding,
   )
@@ -871,8 +894,8 @@ def __init__(
       raise ValueError("Embedding dim for rotary position embedding must be a multiple of 2.")
 
   @property
-  def freqs_cis(self):
-    """Frequencies for rotary embedding."""
+  def corrected_freqs(self):
+    """The per-dimension rotary frequencies, shape [half_dim]."""
     half_dim = self.embedding_dims // 2
     # Compute base frequencies for each (even-indexed) dimension.
     # (Note: We use jnp.arange with float32 for precision.)
@@ -888,15 +911,26 @@ def freqs_cis(self):
     )
     smooth = 1 - self._linear_ramp_factor(low, high, half_dim)
     # The corrected frequency is a weighted mix of the scaled and base values.
-    freqs = freqs / self.rope_factor * (1 - smooth) + freqs * smooth
+    return freqs / self.rope_factor * (1 - smooth) + freqs * smooth
 
+  @property
+  def freqs_cis(self):
+    """Frequencies for every position, shape [max_position_embeddings, half_dim]."""
     # Precompute frequencies for all positions by taking the outer product.
     t = jnp.arange(self.max_position_embeddings, dtype=jnp.float32)  # shape [max_position_embeddings]
     # This gives a [max_position_embeddings, half_dim] tensor with rows as time steps.
-    freqs = jnp.outer(t, freqs)
+    return _cis(jnp.outer(t, self.corrected_freqs))
 
-    # Compute the complex “cis” values: exp(i * theta).
-    return jnp.exp(1j * freqs)  # shape [max_position_embeddings, half_dim]
+  def freqs_cis_at(self, position: Array, out_sharding=None) -> Array:
+    """`freqs_cis[position]`, without building the whole table to index it.
+
+    Row `p` of `freqs_cis` is `exp(1j * p * corrected_freqs)`, so the rows a step
+    actually reads can be produced straight from `position`. The table itself is
+    traced, not a constant XLA can hoist or fold, so building it costs
+    `max_position_embeddings` rows *inside every layer, every step* — on
+    deepseek2-16b that is 163,840 rows of which `max_target_length` are read.
+    """
+    return _cis(jnp.einsum("bs,h->bsh", position.astype(jnp.float32), self.corrected_freqs, out_sharding=out_sharding))
 
   def _find_correction_dim(self, num_rotations: float, dim: int, base: float, max_position_embeddings: int) -> float:
     """Compute the correction dimension for a given number of rotations."""
@@ -967,10 +1001,9 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array:
     else:
       position = position.astype(jnp.int32)
 
-    # Lookup the precomputed frequencies using the position indices.
-    # self.freqs_cis has shape [max_position_embeddings, half_dim] so we use jnp.take along axis 0.
-    # After indexing, shape becomes [B, S, half_dim]; we then add an axis for the heads.
-    freqs = self.freqs_cis.at[position].get(out_sharding=self.freqs_sharding)  # shape: [B, S, half_dim]
+    # Build the frequencies for these positions directly, rather than indexing
+    # them out of the full [max_position_embeddings, half_dim] table.
+    freqs = self.freqs_cis_at(position, out_sharding=self.freqs_sharding)  # shape: [B, S, half_dim]
     freqs = freqs[:, :, jnp.newaxis, :]  # shape: [B, S, 1, half_dim]
 
     if self.interleave and self.pairwise: