tests

Author: wwwjn

torchtitan/models/deepseek_v3/__init__.py

@@ -39,9 +39,9 @@
         n_dense_layers=1,
         n_heads=16,
         moe_args=MoEArgs(
-            num_experts=8,
+            num_experts=2,
             num_shared_experts=2,
-            top_k=3,
+            top_k=1,
             score_func="softmax",
             route_norm=True,
             score_before_experts=False,

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -46,15 +46,15 @@ dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
 
 [parallelism]
 data_parallel_replicate_degree = 1
-data_parallel_shard_degree = -1
+data_parallel_shard_degree = 4
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 1
+tensor_parallel_degree = 2
 enable_async_tensor_parallel = false
 pipeline_parallel_degree = 1
 pipeline_parallel_schedule = "1F1B"
 context_parallel_degree = 1
-expert_parallel_degree = 1
-expert_tensor_parallel_degree = 1
+expert_parallel_degree = 2
+expert_tensor_parallel_degree = 2
 
 [checkpoint]
 enable_checkpoint = false

torchtitan/train.py

@@ -530,6 +530,15 @@ def train_step(
     def train(self):
         job_config = self.job_config
 
+        logger.info("self.parallel_dims: {}".format(self.parallel_dims))
+
+        w1 = self.model_parts[0].layers["1"].moe.experts.w1
+        w2 = self.model_parts[0].layers["1"].moe.experts.w2
+        w3 = self.model_parts[0].layers["1"].moe.experts.w3
+
+        logger.info(f"the GroupedExpert weights w1 shard is: {w1.placements}")
+        logger.info(f"the GroupedExpert weights w2 shard is: {w2.placements}")
+        logger.info(f"the GroupedExpert weights w3 shard is: {w3.placements}")
         self.checkpointer.load(step=job_config.checkpoint.load_step)
         logger.info(f"Training starts at step {self.step + 1}")