test

Author: wwwjn

torchtitan/experiments/kernels/moe/indices.py

@@ -72,13 +72,13 @@ def fill_indices_wrapper(
     write_offsets: torch.Tensor,
     experts_per_rank: int,
     num_ranks: int,
-    total_size: int,
+    max_len: int,
     block_size: int = 128,
-    max_blocks: int = 1024,
+    max_blocks: int = 1024,  # cap on total number of blocks to launch
 ):
-    # Allocate exact size needed instead of max_len
+    # preallocate output
     permuted_indices = torch.full(
-        (total_size,), -1, dtype=torch.int32, device=tokens_per_expert_group.device
+        (max_len,), -1, dtype=torch.int32, device=tokens_per_expert_group.device
     )
 
     # write offsets is per local expert...
@@ -99,37 +99,39 @@ def fill_indices_wrapper(
     return permuted_indices
 
 
-# used for reference testing only
-
-
+# reference
 def fill_indices_cpu(
     tokens_per_expert_group: torch.Tensor,
     start_index_values: torch.Tensor,
     write_offsets: torch.Tensor,
     experts_per_rank: int,
     num_ranks: int,
-    total_size: int,  # Changed from max_len to actual required size
+    max_len: int,
 ):
-    # Allocate exact size needed
+    # We need to preallocate the output - we ignore device and force it on cpu
+    # device = tokens_per_expert_group.device
     permuted_indices = torch.full(
-        (total_size,),
+        (max_len,),
         -1,
         dtype=torch.int32,
-    )
-
+    )  # device=device)
     # Fill the permuted indices
+    # For each local expert
     for e in range(experts_per_rank):
         write_start = write_offsets[e].item()
+        # For each remote rank
         for r in range(num_ranks):
             i = r * experts_per_rank + e
             start_index = start_index_values[i].item()
             length = tokens_per_expert_group[i].item()
+            # Fill in the indices
             if length > 0:
-                end_idx = min(write_start + length, total_size)
+                end_idx = min(write_start + length, max_len)
                 permuted_indices[write_start:end_idx] = torch.arange(
                     start_index,
                     start_index + (end_idx - write_start),
                     dtype=torch.int32,
+                    # device=device,
                 )
             write_start += length
     return permuted_indices
@@ -139,22 +141,24 @@ def generate_permute_indices(
     tokens_per_expert_group: torch.Tensor,
     experts_per_rank: int,
     num_ranks: int,
+    max_len: int,
     alignment: int,
     use_cpu: bool = False,
 ):
     """
     Prepare permutation indices and the number of tokens for each expert.
-    Modified version that returns a tensor of size sum(m_sizes) instead of max_len.
 
     Args:
         tokens_per_expert_group: number of tokens for each expert from all ranks.
         experts_per_rank: number of experts per rank.
         num_ranks: number of ranks.
+        max_len: maximum length of the output index vector.
         alignment: alignment for each returned element in `m_sizes` and padding min for zero token experts.
         use_cpu: whether to use CPU implementation.
 
+
     Returns:
-        permuted_indices: Tensor of indices with size sum(m_sizes), that map original token order to the expert-grouped order.
+        permuted_indices: Tensor of indices that map original token order to the expert-grouped order.
         m_sizes: aligned number of tokens for each expert (padded to alignment boundary).
         m_offsets: Cumulative sum of m_sizes. The exclusive ending position for each expert's tokens.
 
@@ -165,7 +169,7 @@ def generate_permute_indices(
               |  4 |  2 |  1 |  3 |  1 |  2 |  3 |  4 |
     """
 
-    # prefix sum to get start index of each expert
+    # prefix sum to get start index of each expert (parallel scan kernel in future?)
     start_index_values = (
         torch.cumsum(tokens_per_expert_group, 0) - tokens_per_expert_group
     )
@@ -182,12 +186,10 @@ def generate_permute_indices(
     )
 
     # additional prefix sum to get write offset of each expert in permuted_indices
+    # write offsets is per local expert, not global
     m_offsets = torch.cumsum(m_sizes, 0)
     write_offsets = m_offsets - m_sizes
 
-    # Calculate the actual total size needed
-    total_size = m_offsets[-1]
-
     # Select the implementation to use
     if use_cpu:
         permuted_indices = fill_indices_cpu(
@@ -196,16 +198,16 @@ def generate_permute_indices(
             write_offsets,
             experts_per_rank,
             num_ranks,
-            total_size,
+            max_len,
         )
-    else:  # gpu
+    else:
         permuted_indices = fill_indices_wrapper(
             tokens_per_expert_group,
             start_index_values,
             write_offsets,
             experts_per_rank,
             num_ranks,
-            total_size,
+            max_len,
         )
 
     return permuted_indices, m_sizes, m_offsets.to(torch.int32)
@@ -225,17 +227,14 @@ def simple_test():
     alignment = 32
     # Use the GPU kernel
     permuted_indices_gpu, m_sizes, _ = generate_permute_indices(
-        tokens_per_expert_group,
-        experts_per_rank,
-        num_ranks,
-        alignment,
-        use_cpu=False,
+        tokens_per_expert_group, experts_per_rank, num_ranks, max_len, alignment
     )
     # Use the CPU method
     permuted_indices_cpu, m_sizes, _ = generate_permute_indices(
         tokens_per_expert_group,
         experts_per_rank,
         num_ranks,
+        max_len,
         alignment,
         use_cpu=True,
     )
@@ -273,15 +272,16 @@ def test_with_zero_tokens():
         tokens_per_expert_group,
         experts_per_rank,
         num_ranks,
+        max_len,
         alignment,
-        use_cpu=False,
     )
 
     # Use the CPU method
     permuted_indices_cpu, m_sizes_cpu, m_offsets_cpu = generate_permute_indices(
         tokens_per_expert_group,
         experts_per_rank,
         num_ranks,
+        max_len,
         alignment,
         use_cpu=True,
     )

torchtitan/models/deepseek_v3/__init__.py

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

torchtitan/models/deepseek_v3/model/args.py

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

torchtitan/models/deepseek_v3/model/model.py

@@ -14,7 +14,7 @@
 from torchtitan.protocols.train_spec import ModelProtocol
 
 from .args import DeepSeekV3ModelArgs
-from .moe import MoE
+from .moe import FeedForward, MoE
 
 
 # Adapted from https://github.com/DeepSeek-ai/DeepSeek-V3/blob/main/inference/model.py#L294
@@ -260,42 +260,6 @@ def init_weights(self, init_std: float):
             self.q_norm.reset_parameters()
 
 
-class FeedForward(nn.Module):
-    """
-    FeedForward module
-
-    Args:
-        dim (int): Input dimension.
-        hidden_dim (int): Hidden dimension of the feedforward layer.
-        multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
-        ffn_dim_multiplier (float | None): Custom multiplier for hidden dimension. Defaults to None.
-
-    Attributes:
-        w1 (Linear): Linear transformation for the first layer.
-        w2 (Linear): Linear transformation for the second layer.
-        w3 (Linear): Linear transformation for the third layer.
-
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        hidden_dim: int,
-    ):
-        super().__init__()
-        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
-        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
-        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
-
-    def init_weights(self, init_std: float = 0.02):
-        nn.init.trunc_normal_(self.w1.weight, mean=0.0, std=0.02)
-        for linear in (self.w2, self.w3):
-            nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)
-
-
 class TransformerBlock(nn.Module):
     """
     Transformer block with attention and feed-forward layers.
@@ -316,6 +280,7 @@ def __init__(self, layer_id: int, model_args: DeepSeekV3ModelArgs):
 
         # TODO: Need to revisit the weight initialization for the TransformerBlock
         self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5
+        self.layer_id = layer_id
 
     def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):
         """
@@ -330,8 +295,10 @@ 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
 
@@ -360,6 +327,7 @@ 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

@@ -11,6 +11,42 @@
 from .args import DeepSeekV3ModelArgs
 
 
+class FeedForward(nn.Module):
+    """
+    FeedForward module
+
+    Args:
+        dim (int): Input dimension.
+        hidden_dim (int): Hidden dimension of the feedforward layer.
+        multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
+        ffn_dim_multiplier (float | None): Custom multiplier for hidden dimension. Defaults to None.
+
+    Attributes:
+        w1 (Linear): Linear transformation for the first layer.
+        w2 (Linear): Linear transformation for the second layer.
+        w3 (Linear): Linear transformation for the third layer.
+
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+    ):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+    def init_weights(self, init_std: float = 0.02):
+        nn.init.trunc_normal_(self.w1.weight, mean=0.0, std=0.02)
+        for linear in (self.w2, self.w3):
+            nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)
+
+
 # Reference: torchtitan/experiments/llama4/model/
 class GroupedExperts(nn.Module):
     def __init__(
@@ -212,11 +248,17 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
             GroupedExperts(
                 dim=dim,
                 hidden_dim=hidden_dim * model_args.n_shared_experts,
-                num_experts=1,
+                num_experts=1,  # Here needs to be 1 to make it equivalent to the MLP
                 use_grouped_mm=self.use_grouped_mm,
             )
             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
@@ -266,6 +308,15 @@ 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
 
@@ -299,6 +350,7 @@ 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,
                     ALIGN_SIZE_M,
                 )
             token_indices = torch.vstack(
@@ -311,8 +363,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             # NOTE: this would incur a synchronization between device and host
             num_local_tokens_per_expert = num_local_tokens_per_expert.tolist()
 
+        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)
+        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
         )
@@ -321,10 +377,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         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))
 
+        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