tp on groupped_mm finished

Author: wwwjn

torchtitan/experiments/kernels/moe/indices.py

@@ -77,6 +77,7 @@ 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/model/model.py

@@ -295,10 +295,8 @@ def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):
         """
         x = x + self.attention(self.attention_norm(x), freqs_cis)
         if self.moe_enabled:
-            print(f"In TransformerBlock {self.layer_id}: MoE is enabled")
             x = x + self.moe(self.ffn_norm(x))
         else:
-            print(f"In TransformerBlock {self.layer_id}: FFN is enabled")
             x = x + self.feed_forward(self.ffn_norm(x))
         return x
 
@@ -327,7 +325,6 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
 
         self.layers = torch.nn.ModuleDict()
         for layer_id in range(model_args.n_layers):
-            print(f"Create layer: {layer_id}")
             self.layers[str(layer_id)] = TransformerBlock(layer_id, model_args)
 
         self.norm = nn.RMSNorm(model_args.dim)

torchtitan/models/deepseek_v3/model/moe.py

@@ -211,7 +211,7 @@ def forward(
         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):
@@ -253,12 +253,6 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
             )
             if model_args.n_shared_experts > 0
             else None
-            # FeedForward(
-            #     dim=dim,
-            #     hidden_dim=hidden_dim * model_args.n_shared_experts,
-            # )
-            # if model_args.n_shared_experts > 0
-            # else None
         )
 
         # auxiliary-loss-free load balancing
@@ -298,6 +292,7 @@ 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,)
@@ -308,14 +303,14 @@ 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,
-        )
+        # 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
@@ -329,6 +324,12 @@ 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
@@ -350,28 +351,31 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                     num_local_tokens_per_expert,
                     self.experts.num_experts,
                     1,
-                    token_indices[0] + self.experts.num_experts * ALIGN_SIZE_M,
+                    token_indices.shape[0] + self.experts.num_experts * ALIGN_SIZE_M,
                     ALIGN_SIZE_M,
                 )
-            token_indices = torch.vstack(
-                (token_indices, token_indices.new_zeros((dim)))
-            )
-            token_indices = token_indices[permuted_indices, :]
+
             routed_input = torch.vstack((routed_input, routed_input.new_zeros((dim))))
+            input_shape = routed_input.shape
             routed_input = routed_input[permuted_indices, :]
         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
 
-        print("Num local tokens per expert: ", num_local_tokens_per_expert)
         # shape (bs*slen*top_k, dim)
         routed_output = self.experts(
             routed_input, num_local_tokens_per_expert
         )  # torch.Size([16384(bsz), 256])
-        print("Routed output shape: ", routed_output.shape)
-        routed_output = (routed_output.to(torch.float32) * top_scores.unsqueeze(-1)).to(
-            x.dtype
-        )
+
+        routed_output_unpermuted = routed_output.new_empty(input_shape)
+        routed_output_unpermuted[permuted_indices, :] = routed_output
+        routed_output = routed_output_unpermuted[:-1]
+
+        # 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
+        # )
 
         # shared expert
         if self.shared_expert is not None:
@@ -381,10 +385,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         else:
             out = torch.zeros_like(x.reshape(bs * slen, dim))
 
-        print(
-            "Out shape: ", out.shape, out.grad.shape if out.grad is not None else None
-        )
-
         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/debug_model.toml

@@ -45,6 +45,7 @@ max_norm = 1.0  # grad norm clipping
 steps = 10
 compile = false
 dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
+seed = 0
 
 [parallelism]
 data_parallel_replicate_degree = 1

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -41,15 +41,15 @@ lr_min = 0.0
 local_batch_size = 16
 seq_len = 2048
 max_norm = 1.0  # grad norm clipping
-steps = 10
+steps = 2
 compile = false
 dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
 
 [parallelism]
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 1
+tensor_parallel_degree = 2
 enable_async_tensor_parallel = false
 
 [checkpoint]