Add support for pipeling deepseek

Author: gobbleturk

MaxText/layers/models.py

@@ -208,7 +208,7 @@ def setup(self):
     self.decoder_layer = self.get_decoder_layers()
     self.norm_layer = self.get_norm_layer()
     if self.config.using_pipeline_parallelism:
-      pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer[0])
+      pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer)
       remat_policy = self.get_remat_policy()
       self.pipeline_module = pipeline.Pipeline(
           config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy
@@ -397,8 +397,15 @@ def scan_decoder_layers(self, cfg, decoder_layer, length, metdata_axis_name, mes
     )
     return scan_fn(config=cfg, mesh=mesh, name=metdata_axis_name, quant=self.quant)
 
-  def get_pipeline_stage_module(self, base_stage):
+  def get_pipeline_stage_module(self, decoder_blocks):
+    def get_layer_to_pipeline(blocks, cfg):
+      if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
+        return blocks[1] # return the sparse block
+      else:
+        return blocks[0]
     cfg = self.config
+    base_stage = get_layer_to_pipeline(decoder_blocks, cfg)
+
     if cfg.set_remat_policy_on_layers_per_stage:
       policy = self.get_remat_policy()
       base_stage = self.set_remat_policy([base_stage], policy)[0]
@@ -463,20 +470,46 @@ def __call__(
         )
       else:
         partition_spec = None  # This partition spec is only used for the fsdp_ag_once feature.
-      y = self.pipeline_module(
-          y, decoder_segment_ids, decoder_positions, deterministic, model_mode, partition_spec=partition_spec
-      )
-      remaining_layers = self.config.num_decoder_layers - self.config.pipeline_parallel_layers
-      if remaining_layers > 0:
+      if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
+        assert len(RemattedBlockLayers) == 2, f"Scanned layers must have a length of 2 using deepseek."
+        dense_layer = RemattedBlockLayers[0]
+        moe_layer = RemattedBlockLayers[1]
+        num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
+        num_moe_layers_outside_pp = num_moe_layers - self.config.pipeline_parallel_layers
         logical_axis_rules_pp_as_dp = maxtext_utils.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules)
+        # We chose not to pipeline the dense layers, only sparse for SPMD.
         with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp):
-          y, _ = self.scan_decoder_layers(cfg, RemattedBlockLayers[0], remaining_layers, "layers", mesh)(
+          y, _ = self.scan_decoder_layers(cfg, dense_layer, cfg.first_num_dense_layers, "dense_layers", mesh)(
               y,
               decoder_segment_ids,
               decoder_positions,
               deterministic,
               model_mode,
           )
+          if num_moe_layers_outside_pp > 0:
+            y, _ = self.scan_decoder_layers(cfg, moe_layer, num_moe_layers_outside_pp, "moe_layers", mesh)(
+                y,
+                decoder_segment_ids,
+                decoder_positions,
+                deterministic,
+                model_mode,
+            )
+        y = self.pipeline_module(y, decoder_segment_ids, decoder_positions, deterministic, model_mode, partition_spec=partition_spec)
+      else:
+        y = self.pipeline_module(
+            y, decoder_segment_ids, decoder_positions, deterministic, model_mode, partition_spec=partition_spec
+        )
+        remaining_layers = self.config.num_decoder_layers - self.config.pipeline_parallel_layers
+        if remaining_layers > 0:
+          logical_axis_rules_pp_as_dp = maxtext_utils.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules)
+          with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp):
+            y, _ = self.scan_decoder_layers(cfg, RemattedBlockLayers[0], remaining_layers, "layers", mesh)(
+                y,
+                decoder_segment_ids,
+                decoder_positions,
+                deterministic,
+                model_mode,
+            )
     else:
       if cfg.scan_layers:
         if cfg.decoder_block == DecoderBlockType.DEEPSEEK:

MaxText/pyconfig.py

@@ -850,8 +850,6 @@ def pipeline_first_axis(raw_keys):
 
 
 def validate_deepseek_moe(raw_keys):
-  if raw_keys["decoder_block"] == "deepseek" and using_pipeline_parallelism(raw_keys):
-    raise ValueError("Currently we do not support DeepSeek MoE with pipeline parallelism.")
   if raw_keys["n_routing_groups"] != -1:
     if raw_keys["topk_routing_group"] == -1:
       raise ValueError(f'config topk_routing_group: {raw_keys["topk_routing_group"]} is not defined')

MaxText/tests/train_compile_test.py

@@ -559,6 +559,30 @@ def test_moe_deepseek_without_device_limit(self):
         )
     )
 
+
+  @pytest.mark.tpu_only
+  def test_moe_deepseek_pipeline_subset(self):
+    compiled_trainstep_file = "/tmp/test_moe_deepseek_pipeline_subset.pickle"
+    train_compile_main(
+        (
+            None,
+            os.path.join(PKG_DIR, "configs", "base.yml"),
+            f"compiled_trainstep_file={compiled_trainstep_file}",
+            "compile_topology=v6e-256",
+            "compile_topology_num_slices=8",
+            "use_iota_embed=true",
+            "model_name=deepseek3-671b",
+            "sparse_matmul=False",
+            "megablox=False",
+            "capacity_factor=1",
+            "per_device_batch_size=1",
+            "max_target_length=2048",
+            "pipeline_parallel_layers=56",
+            "ici_expert_parallelism=16",
+            "dcn_pipeline_parallelism=8"
+        )
+    )
+
   @pytest.mark.skip(reason="b/415132665: Enable it once scan is supported in training for shorter compiler time")
   @pytest.mark.tpu_only
   def test_moe_llama4_17b_16e(self):