test

Author: wwwjn

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

@@ -76,7 +76,7 @@ class DeepSeekV3ModelArgs(BaseModelArgs):
     score_func: Literal["softmax", "sigmoid"] = "softmax"
     route_scale: float = 1.0
     use_grouped_mm: bool = False
-    load_balance_coeff: float | None = 1e-3
+    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
 
@@ -311,8 +362,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 +376,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

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -42,7 +42,7 @@ lr_min = 0.0
 local_batch_size = 8
 seq_len = 2048
 max_norm = 1.0  # grad norm clipping
-steps = 10
+steps = 2
 compile = false
 dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
 
@@ -62,7 +62,7 @@ export_dtype = "float32"
 async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem"]
 
 [activation_checkpoint]
-mode = "none"  # ["none", "selective", "full"]
+mode = "full"  # ["none", "selective", "full"]
 selective_ac_option = '2'  # 'int' = ac every positive int layer or 'op', ac based on ops policy
 
 [float8]

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -49,7 +49,7 @@ dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
 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]

torchtitan/train.py

@@ -11,10 +11,10 @@
 from typing import Any, Generator, Iterable, Optional
 
 import torch
-from torch.distributed.elastic.multiprocessing.errors import record
 
 import torchtitan.components.ft as ft
 import torchtitan.protocols.train_spec as train_spec_module
+from torch.distributed.elastic.multiprocessing.errors import record
 from torchtitan.components.checkpoint import CheckpointManager
 from torchtitan.components.dataloader import DataloaderStopIteration
 from torchtitan.components.loss import rescale_accumulated_loss
@@ -551,5 +551,5 @@ def close(self) -> None:
             trainer.close()
 
         if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
+            # torch.distributed.destroy_process_group()
             logger.info("Process group destroyed.")