support deepseek explicit shard mode

Author: NuojCheng

src/MaxText/common_types.py

@@ -37,7 +37,9 @@
 PREFILL_LENGTH = "prefill_activation_length"
 Q_LENGTH = "activation_q_length"
 Q_LENGTH_NO_EXP = "activation_q_length_no_exp"
+Q_LORA_UP_PROJ = "q_lora_up_proj"
 KV_LENGTH = "activation_kv_length"
+KV_LORA_UP_PROJ = "kv_lora_up_proj"
 EMBED = "activation_embed"
 HEAD = "activation_heads"
 PREFILL_KV_BATCH = "activation_prefill_kv_batch"

src/MaxText/configs/types.py

@@ -1848,7 +1848,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         if self.packing:
           raise ValueError("For multimodal SFT, `packing` is not yet supported.")
     if self.shard_mode == ShardMode.EXPLICIT:
-      supported_decoders = {"simple", "simple_mlp", "llama2"}
+      supported_decoders = {"simple", "simple_mlp", "llama2", "deepseek"}
       if self.decoder_block.value not in supported_decoders:
         raise ValueError(
             f"Decoder '{self.decoder_block.value}' is not supported with 'explicit' sharding. "

src/MaxText/gradient_accumulation.py

@@ -92,6 +92,7 @@ def convert_to_bf16(param):
   def accumulate_gradient(acc_grad_and_loss, data):
     ga_params = acc_grad_and_loss["ga_params"]
     (_, aux), cur_batch_gradient = grad_func(model, config, data, dropout_rng, ga_params, *extra_dpo_args, is_train=True)
+    cur_batch_gradient = jax.tree.map(_maybe_shard_with_name, cur_batch_gradient, grad_shardings)
     acc_grad_and_loss["loss"] += aux["total_loss"]
     acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"]
     acc_grad_and_loss["mtp_loss"] += aux["mtp_loss"]

src/MaxText/layers/attention_mla.py

@@ -16,8 +16,10 @@
 # pylint: disable=arguments-differ
 # pylint: disable=no-name-in-module
 
+from functools import partial
+
 from jax.ad_checkpoint import checkpoint_name
-from jax.sharding import Mesh
+from jax.sharding import Mesh, NamedSharding
 import jax.numpy as jnp
 
 from flax import linen as nn
@@ -31,6 +33,7 @@
 from MaxText.layers import quantizations
 from MaxText.layers.quantizations import AqtQuantization as Quant
 from MaxText.inference import page_manager
+from MaxText.sharding import maybe_shard_with_logical
 from MaxText.common_types import MODEL_MODE_PREFILL
 
 # -----------------------------------------
@@ -66,8 +69,14 @@ def self_attention_with_norm(
     logical_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_embed")
   else:
     logical_axis_names = ("activation_batch", "activation_norm_length", "activation_embed")
+  _maybe_shard_with_logical = partial(
+      maybe_shard_with_logical,
+      mesh=mesh,
+      shard_mode=cfg.shard_mode,
+  )
+  lnx_out_sharding = NamedSharding(mesh, nn.logical_to_mesh_axes(logical_axis_names))
 
-  lnx = nn.with_logical_constraint(lnx, logical_axis_names)
+  lnx = _maybe_shard_with_logical(lnx, logical_axis_names)
 
   attention_layer = attention_mla.mla_as_linen(
       config=cfg,
@@ -106,12 +115,13 @@ def self_attention_with_norm(
       decoder_segment_ids=decoder_segment_ids,
       deterministic=deterministic,
       model_mode=model_mode,
+      out_sharding=lnx_out_sharding,
       previous_chunk=previous_chunk,
       page_state=page_state,
       slot=slot,
   )
 
-  attention_lnx = nn.with_logical_constraint(attention_lnx, logical_axis_names)
+  attention_lnx = _maybe_shard_with_logical(attention_lnx, logical_axis_names)
   intermediate_inputs = inputs + attention_lnx
 
   # Normalization
@@ -123,7 +133,7 @@ def self_attention_with_norm(
       kernel_axes=("norm",),
       epsilon=cfg.normalization_layer_epsilon,
   )(intermediate_inputs)
-  hidden_states = nn.with_logical_constraint(hidden_states, logical_axis_names)
+  hidden_states = _maybe_shard_with_logical(hidden_states, logical_axis_names)
   return hidden_states, intermediate_inputs
 
 
@@ -167,9 +177,14 @@ def __call__(
     cfg = self.config
     if model_mode == MODEL_MODE_PREFILL:
       logical_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_embed")
+      mlp_logical_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_mlp")
     else:
       logical_axis_names = ("activation_batch", "activation_norm_length", "activation_embed")
-    inputs = nn.with_logical_constraint(inputs, logical_axis_names)
+      mlp_logical_axis_names = ("activation_batch", "activation_norm_length", "activation_mlp")
+    _maybe_shard_with_logical = partial(maybe_shard_with_logical, mesh=self.mesh, shard_mode=self.config.shard_mode)
+    lnx_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(logical_axis_names))
+    mlp_intermediate_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(mlp_logical_axis_names))
+    inputs = _maybe_shard_with_logical(inputs, logical_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
 
     hidden_states, intermediate_inputs = self_attention_with_norm(
@@ -196,12 +211,17 @@ def __call__(
         config=cfg,
         mesh=self.mesh,
         quant=self.quant,
-    )(hidden_states, deterministic=deterministic)
-    mlp_lnx = nn.with_logical_constraint(mlp_lnx, logical_axis_names)
+    )(
+        hidden_states,
+        deterministic=deterministic,
+        intermediate_sharding=mlp_intermediate_sharding,
+        out_sharding=lnx_out_sharding,
+    )
+    mlp_lnx = _maybe_shard_with_logical(mlp_lnx, logical_axis_names)
 
     layer_output = mlp_lnx + intermediate_inputs
     layer_output = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(layer_output, deterministic=deterministic)
-    layer_output = nn.with_logical_constraint(
+    layer_output = _maybe_shard_with_logical(
         layer_output,
         logical_axis_names,
     )
@@ -234,9 +254,14 @@ def __call__(
     cfg = self.config
     if model_mode == MODEL_MODE_PREFILL:
       logical_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_embed")
+      mlp_logical_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_mlp")
     else:
       logical_axis_names = ("activation_batch", "activation_norm_length", "activation_embed")
-    inputs = nn.with_logical_constraint(inputs, logical_axis_names)
+      mlp_logical_axis_names = ("activation_batch", "activation_norm_length", "activation_mlp")
+    _maybe_shard_with_logical = partial(maybe_shard_with_logical, mesh=self.mesh, shard_mode=self.config.shard_mode)
+    lnx_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(logical_axis_names))
+    lnx_intermediate_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(mlp_logical_axis_names))
+    inputs = _maybe_shard_with_logical(inputs, logical_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
 
     hidden_states, intermediate_inputs = self_attention_with_norm(
@@ -265,12 +290,12 @@ def __call__(
         dtype=cfg.dtype,
         weight_dtype=cfg.weight_dtype,
         quant=self.quant,
-    )(hidden_states)
-    mlp_lnx = nn.with_logical_constraint(mlp_lnx, logical_axis_names)
+    )(hidden_states, intermediate_sharding=lnx_intermediate_sharding, out_sharding=lnx_out_sharding)
+    mlp_lnx = _maybe_shard_with_logical(mlp_lnx, logical_axis_names)
 
     layer_output = mlp_lnx + intermediate_inputs
     layer_output = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(layer_output, deterministic=deterministic)
-    layer_output = nn.with_logical_constraint(
+    layer_output = _maybe_shard_with_logical(
         layer_output,
         logical_axis_names,
     )

src/MaxText/layers/attentions.py

@@ -263,9 +263,11 @@ def __init__(
       self,
       min_timescale: int,
       max_timescale: int,
+      mesh: Mesh,
       embedding_dims: int = 0,
       cast_as_fprop_dtype: bool = True,
       fprop_dtype: DType = jnp.bfloat16,
+      shard_mode: ShardMode = ShardMode.AUTO,
       # Not used in RotaryEmbedding but passed in by nnx.bridge.to_linen.
       # TODO: Remove when bridge no longer needed
       rope_linear_scaling_factor: float = 1.0,
@@ -285,9 +287,11 @@ def __init__(
     """
     self.min_timescale = min_timescale
     self.max_timescale = max_timescale
+    self.mesh = mesh
     self.embedding_dims = embedding_dims
     self.cast_as_fprop_dtype = cast_as_fprop_dtype
     self.fprop_dtype = fprop_dtype
+    self.shard_mode = shard_mode
     self.rope_linear_scaling_factor = rope_linear_scaling_factor
 
     if self.embedding_dims % 2:
@@ -384,6 +388,7 @@ def qwen3_next_rotary_embedding_as_linen(
     *,
     min_timescale: int,
     max_timescale: int,
+    mesh: Mesh,
     embedding_dims: int = 0,
     partial_rotary_factor: float = 0.25,
     cast_as_fprop_dtype: bool = True,
@@ -407,6 +412,7 @@ def qwen3_next_rotary_embedding_as_linen(
       Qwen3NextRotaryEmbedding,
       min_timescale=min_timescale,
       max_timescale=max_timescale,
+      mesh=mesh,
       embedding_dims=embedding_dims,
       partial_rotary_factor=partial_rotary_factor,
       cast_as_fprop_dtype=cast_as_fprop_dtype,
@@ -423,6 +429,7 @@ def __init__(
       self,
       min_timescale: int,
       max_timescale: int,
+      mesh: Mesh,
       embedding_dims: int = 0,
       cast_as_fprop_dtype: bool = True,
       fprop_dtype: DType = jnp.bfloat16,
@@ -447,6 +454,7 @@ def __init__(
     super().__init__(
         min_timescale=min_timescale,
         max_timescale=max_timescale,
+        mesh=mesh,
         embedding_dims=self.rotary_dim,
         cast_as_fprop_dtype=cast_as_fprop_dtype,
         fprop_dtype=fprop_dtype,
@@ -478,10 +486,12 @@ def __init__(
       self,
       min_timescale: int,
       max_timescale: int,
+      mesh: Mesh,
       embedding_dims: int = 0,
       cast_as_fprop_dtype: bool = True,
       fprop_dtype: DType = jnp.bfloat16,
       use_scale: bool = True,
+      shard_mode: ShardMode = ShardMode.AUTO,
       # Not used in LLaMARotaryEmbedding but passed in by nnx.bridge.to_linen.
       # TODO: Remove when bridge no longer needed
       rngs: nnx.Rngs = None,
@@ -505,6 +515,8 @@ def __init__(
         embedding_dims=embedding_dims,
         cast_as_fprop_dtype=cast_as_fprop_dtype,
         fprop_dtype=fprop_dtype,
+        mesh=mesh,
+        shard_mode=shard_mode,
         rngs=rngs,
     )
 
@@ -613,6 +625,7 @@ def __call__(self, inputs: jax.Array, position: None | jax.Array = None) -> jax.
 def yarn_rotary_embedding_as_linen(
     *,
     embedding_dims: int,
+    mesh: Mesh,
     max_position_embeddings: int = 4096 * 4,
     original_max_position_embeddings: int = 4096,
     beta_fast: float = 32,
@@ -625,6 +638,7 @@ def yarn_rotary_embedding_as_linen(
     interleave: bool = True,
     truncate: bool = True,
     attention_scaling: bool = False,
+    shard_mode: ShardMode = ShardMode.AUTO,
 ):
   """Initializes the YarnRotaryEmbedding module and returns it as a Linen module.
 
@@ -643,6 +657,7 @@ def yarn_rotary_embedding_as_linen(
   return nnx_wrappers.to_linen(
       YarnRotaryEmbedding,
       embedding_dims=embedding_dims,
+      mesh=mesh,
       max_position_embeddings=max_position_embeddings,
       original_max_position_embeddings=original_max_position_embeddings,
       beta_fast=beta_fast,
@@ -656,6 +671,7 @@ def yarn_rotary_embedding_as_linen(
       interleave=interleave,
       truncate=truncate,
       attention_scaling=attention_scaling,
+      shard_mode=shard_mode,
   )
 
 
@@ -685,6 +701,7 @@ class YarnRotaryEmbedding(nnx.Module):
   def __init__(
       self,
       embedding_dims: int,
+      mesh: Mesh,
       max_position_embeddings: int = 4096 * 4,
       original_max_position_embeddings: int = 4096,
       beta_fast: float = 32,
@@ -693,6 +710,7 @@ def __init__(
       rope_factor: float = 40,
       cast_as_fprop_dtype: bool = True,
       fprop_dtype: DType = jnp.bfloat16,
+      shard_mode: ShardMode = ShardMode.AUTO,
       interleave=True,
       truncate=True,
       attention_scaling=False,
@@ -712,6 +730,13 @@ def __init__(
     self.fprop_dtype = fprop_dtype
     self.interleave = interleave
     self.truncate = truncate
+    self.mesh = mesh
+    self.shard_mode = shard_mode
+    self.freqs_sharding = (
+        NamedSharding(mesh, nn.logical_to_mesh_axes(("activation_batch", "activation_length_no_exp", "q_heads")))
+        if shard_mode == ShardMode.EXPLICIT
+        else None
+    )
     self.attention_scaling = attention_scaling
 
     if self.embedding_dims % 2:
@@ -811,7 +836,8 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array:
     # 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 = jnp.take(self.freqs_cis, position, axis=0)  # shape: [B, S, half_dim]
+    # freqs = jnp.take(self.freqs_cis, position, axis=0)  # shape: [B, S, half_dim]
+    freqs = self.freqs_cis.at[position].get(out_sharding=self.freqs_sharding)
     freqs = freqs[:, :, jnp.newaxis, :]  # shape: [B, S, 1, half_dim]
 
     if self.interleave:
@@ -828,7 +854,14 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array:
 
     inputs_complex = first_half + 1j * second_half  # shape: [B, S, N, half_dim]
     # Apply the rotary transformation via complex multiplication.
-    rotated = inputs_complex * freqs  # shape: [B, S, N, half_dim]
+    rotated_sharding = (
+        NamedSharding(self.mesh, nn.logical_to_mesh_axes(("activation_batch", "activation_length_no_exp", None, None)))
+        if self.shard_mode == ShardMode.EXPLICIT
+        else None
+    )
+    rotated = jnp.einsum(
+        "ijkl, ijml->ijkl", inputs_complex, freqs, out_sharding=rotated_sharding
+    )  # shape: [B, S, N, half_dim]
     # Convert the complex result back to a real tensor.
     # Split the complex number into its real and imaginary parts.
     # [real1, real2, ..., img1, img2, ...]