refactor pipeline code

Author: NuojCheng

src/maxtext/configs/types.py

@@ -2526,39 +2526,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.")
 
     # I. FINAL TYPE CONVERSIONS AND DERIVED LISTS
-    # Create the ici_parallelism and dcn_parallelism lists for legacy compatibility.
-    # if self.using_pipeline_parallelism and self.mesh_axes and self.mesh_axes[0] == "stage":
-    #   self.ici_parallelism = [
-    #       self.ici_diloco_parallelism,
-    #       self.ici_pipeline_parallelism,
-    #       self.ici_data_parallelism,
-    #       self.ici_fsdp_parallelism,
-    #       self.ici_fsdp_transpose_parallelism,
-    #       self.ici_sequence_parallelism,
-    #       self.ici_context_parallelism,
-    #       self.ici_context_autoregressive_parallelism,
-    #       self.ici_tensor_parallelism,
-    #       self.ici_tensor_transpose_parallelism,
-    #       self.ici_tensor_sequence_parallelism,
-    #       self.ici_expert_parallelism,
-    #       self.ici_autoregressive_parallelism,
-    #   ]
-    #   self.dcn_parallelism = [
-    #       self.dcn_diloco_parallelism,
-    #       self.dcn_pipeline_parallelism,
-    #       self.dcn_data_parallelism,
-    #       self.dcn_fsdp_parallelism,
-    #       self.dcn_fsdp_transpose_parallelism,
-    #       self.dcn_sequence_parallelism,
-    #       self.dcn_context_parallelism,
-    #       self.dcn_context_autoregressive_parallelism,
-    #       self.dcn_tensor_parallelism,
-    #       self.dcn_tensor_transpose_parallelism,
-    #       self.dcn_tensor_sequence_parallelism,
-    #       self.dcn_expert_parallelism,
-    #       self.dcn_autoregressive_parallelism,
-    #   ]
-    # else:
     ici_map = {
         "diloco": self.ici_diloco_parallelism,
         "data": self.ici_data_parallelism,

src/maxtext/layers/attention_op.py

@@ -580,9 +580,7 @@ def maybe_create_nnx(einsum, *args):
 
   def _logical_to_mesh_axes(self, logical_name):
     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
-    )
+    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

@@ -307,7 +307,7 @@ def setup(self):
     if self.config.using_pipeline_parallelism:
       pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer)
       remat_policy = self.get_remat_policy()
-      self.pipeline_module = pipeline.Pipeline(
+      self.pipeline_module = pipeline.create_pipeline(
           config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy
       )
 
@@ -794,8 +794,10 @@ def __call__(
         model_mode,
     )
     if cfg.using_pipeline_parallelism:
-      logical_partition_spec = self.pipeline_module.get_weight_sharding(
-          y, decoder_segment_ids, decoder_positions, deterministic, model_mode
+      logical_partition_spec = (
+          self.pipeline_module.get_weight_sharding(y, decoder_segment_ids, decoder_positions, deterministic, model_mode)
+          if cfg.quantization == ""
+          else None
       )
       if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
         assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek."