TP gemm works

Author: wwwjn

torchtitan/experiments/kernels/moe/indices.py

@@ -77,7 +77,6 @@ def fill_indices_wrapper(
     max_blocks: int = 1024,  # cap on total number of blocks to launch
 ):
     # preallocate output
-    print("max_len: ", max_len, "block_size: ", block_size, "max_blocks: ", max_blocks)
     permuted_indices = torch.full(
         (max_len,), -1, dtype=torch.int32, device=tokens_per_expert_group.device
     )

torchtitan/models/deepseek_v3/__init__.py

@@ -32,12 +32,12 @@
         dim=256,
         inter_dim=10944,
         moe_inter_dim=1408,
-        n_layers=1,
-        n_dense_layers=0,  # no FFN layer, all MoE layers
+        n_layers=3,
+        n_dense_layers=1,
         n_heads=16,
-        n_routed_experts=2,  # hang only happens when n_routed_experts > n_activated_experts
-        n_shared_experts=1,
-        n_activated_experts=1,
+        n_routed_experts=8,
+        n_shared_experts=2,
+        n_activated_experts=3,
         route_scale=1.0,
         q_lora_rank=0,
         kv_lora_rank=512,

torchtitan/models/deepseek_v3/model/args.py

@@ -75,7 +75,7 @@ class DeepSeekV3ModelArgs(BaseModelArgs):
     n_limited_groups: int = 1
     score_func: Literal["softmax", "sigmoid"] = "softmax"
     route_scale: float = 1.0
-    use_grouped_mm: bool = True
+    use_grouped_mm: bool = False
     load_balance_coeff: float = 1e-3
     # Multi-Head Latent Attention (MLA)
     q_lora_rank: int = 0

torchtitan/models/deepseek_v3/model/moe.py

@@ -201,17 +201,20 @@ def forward(
             min=0,
             max=self.num_experts,
         )
+
+        # Reorder the token indices to match the order of the experts
         # token_indices_experts_sorted shape (bs*slen*top_k,)
         token_indices_experts_sorted = torch.argsort(
             selected_experts_indices.view(-1), stable=True
         )
+
+        # reorder the scores to match the order of the token indices
         top_scores = top_scores.view(-1)[token_indices_experts_sorted]
         token_indices_experts_sorted = token_indices_experts_sorted // self.top_k
 
         top_scores = (
             top_scores * self.route_sclaing_factor
         )  # must multiply the scaling factor
-        print("In TokenChoiceTopKRouter, top_scores shape: ", top_scores)
         return top_scores, token_indices_experts_sorted, num_local_tokens_per_expert
 
     def init_weights(self, init_std: float):
@@ -292,7 +295,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         Returns:
             out (torch.Tensor): Output tensor with shape ``(bs, slen, dim)``.
         """
-        print("In MoE input, x shape: ", x)
         bs, slen, dim = x.shape
 
         # top_scores and selected_indices shape (bs*slen*top_k,)
@@ -303,15 +305,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             num_local_tokens_per_expert,
         ) = self.router(x.reshape(bs * slen, dim), self.expert_bias)
 
-        # print(
-        #     "In MoE, top_scores shape: ",
-        #     top_scores.shape,
-        #     "token_indices: ",
-        #     token_indices.shape,
-        #     "num_local_tokens: ",
-        #     num_local_tokens_per_expert.shape,
-        # )
-
         # will be used to update the expert bias for load balancing
         self.tokens_per_expert += num_local_tokens_per_expert
 
@@ -324,12 +317,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             dim=0,
             index=token_indices,
         )
-        print("Routed input: ", routed_input)
-
-        # TODO: remove this line, this is a temporary test
-        routed_input = (routed_input.to(torch.float32) * top_scores.reshape(-1, 1)).to(
-            x.dtype
-        )
 
         if self.use_grouped_mm:
             # NOTE: In order to use torch._grouped_mm, we need to make sure
@@ -361,30 +348,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         else:
             # NOTE: this would incur a synchronization between device and host
             num_local_tokens_per_expert = num_local_tokens_per_expert.tolist()
-            input_shape, permuted_indices = None, None
+            permuted_indices, input_shape = None, None
 
         # shape (bs*slen*top_k, dim)
-        routed_output = self.experts(
-            routed_input, num_local_tokens_per_expert
-        )  # torch.Size([16384(bsz), 256])
+        routed_output = self.experts(routed_input, num_local_tokens_per_expert)
 
-        routed_output_unpermuted = routed_output.new_empty(input_shape)
-        routed_output_unpermuted[permuted_indices, :] = routed_output
-        routed_output = routed_output_unpermuted[:-1]
+        if self.use_grouped_mm:
+            # NOTE: Reverese the permutation to get the original order as inputs
+            routed_output_unpermuted = routed_output.new_empty(input_shape)
+            routed_output_unpermuted[permuted_indices, :] = routed_output
+            routed_output = routed_output_unpermuted[:-1]  # remove padding
 
-        # TODO: Use this line instead if routed_input*top_scores, need to pad top_scores to be multiple of 16
-        # routed_output = (routed_output.to(torch.float32) * top_scores.unsqueeze(-1)).to(
-        #     x.dtype
-        # )
+        routed_output = (routed_output.to(torch.float32) * top_scores.unsqueeze(-1)).to(
+            x.dtype
+        )
 
         # shared expert
         if self.shared_expert is not None:
             out = self.shared_expert(x.reshape(1, bs * slen, dim)).reshape(
                 bs * slen, dim
-            )  #  torch.Size([16384, 256]) None
+            )
         else:
             out = torch.zeros_like(x.reshape(bs * slen, dim))
 
+        # Accumulate multiple expert results becase each token can be routed to multiple experts
         out = out.scatter_add(dim=0, index=token_indices, src=routed_output)
         out = out.reshape(bs, slen, dim)
         return out

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -15,7 +15,7 @@ save_memory_snapshot_folder = "memory_snapshot"
 [metrics]
 log_freq = 1
 disable_color_printing = false
-enable_tensorboard = false
+enable_tensorboard = true
 save_tb_folder = "tb"
 enable_wandb = false
 
@@ -41,15 +41,16 @@ lr_min = 0.0
 local_batch_size = 16
 seq_len = 2048
 max_norm = 1.0  # grad norm clipping
-steps = 2
+steps = 10
 compile = false
 dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
+seed = 0
 
 [parallelism]
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 2
+tensor_parallel_degree = 1
 enable_async_tensor_parallel = false
 
 [checkpoint]