enable pp with batch split ds

Author: NuojCheng

src/MaxText/layers/attention_op.py

@@ -524,7 +524,8 @@ def maybe_create_nnx(einsum, *args):
       self.AqtEinsum_3 = jnp.einsum
 
   def _logical_to_mesh_axes(self, logical_name):
-    return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=self.config.logical_axis_rules)
+    logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules
+    return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules)
 
   def check_attention_inputs(self, query: Array, key: Array | KVTensor, value: Array | KVTensor) -> None:
     """Check attention inputs."""

src/MaxText/layers/decoders.py

@@ -954,11 +954,11 @@ def __call__(
     else:
       logits = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode)
       logits = sharding.maybe_shard_with_logical(
-        logits,
-        ("activation_embed_and_logits_batch", "activation_length_no_exp", "activation_vocab"),
-        mesh=self.mesh,
-        shard_mode=self.config.shard_mode,
-        debug_sharding=self.config.debug_sharding,
+          logits,
+          ("activation_embed_and_logits_batch", "activation_length_no_exp", "activation_vocab"),
+          mesh=self.mesh,
+          shard_mode=self.config.shard_mode,
+          debug_sharding=self.config.debug_sharding,
       )
 
     # The API of the Decoder is now a tuple, providing both the main output

src/MaxText/layers/deepseek_batchsplit.py

@@ -755,7 +755,7 @@ def gmm(
       input_buffer_count,
       combine_scopes,
   ):
-    if config.use_qwix_quantization:
+    if config.use_qwix_quantization or config.using_pipeline_parallelism:
       output = megablox.gmm(
           lhs=inputs,
           rhs=kernel,

src/MaxText/layers/moe.py

@@ -462,7 +462,8 @@ def _maybe_shard_with_logical(self, inputs, logical_name):
     )
 
   def _logical_to_mesh_axes(self, logical_name):
-    return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=self.config.logical_axis_rules)
+    logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules
+    return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules)
 
   def get_expert_parallelism_size(self):
     return self.mesh.shape.get("expert", 1)