fix bug

Author: wwwjn

torchtitan/models/deepseek_v3/model/model.py

@@ -152,17 +152,23 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
         self.v_head_dim = model_args.v_head_dim
 
         if self.q_lora_rank == 0:
-            self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim)
+            self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim, bias=False)
         else:
-            self.wq_a = nn.Linear(self.dim, self.q_lora_rank)
+            self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False)
             self.q_norm = nn.RMSNorm(self.q_lora_rank, eps=model_args.norm_eps)
-            self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim)
-        self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)
+            self.wq_b = nn.Linear(
+                self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False
+            )
+        self.wkv_a = nn.Linear(
+            self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False
+        )
         self.kv_norm = nn.RMSNorm(self.kv_lora_rank, eps=model_args.norm_eps)
         self.wkv_b = nn.Linear(
-            self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim)
+            self.kv_lora_rank,
+            self.n_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
         )
-        self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim)
+        self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False)
         self.softmax_scale = self.qk_head_dim**-0.5
 
         if model_args.max_seq_len > model_args.original_seq_len:
@@ -192,8 +198,8 @@ def forward(
         if self.q_lora_rank == 0:
             q = self.wq(x)  # (bsz, seqlen, n_heads * qk_head_dim)
         else:
-            q = self.wq_b(self.q_norm(self.wq_a(x)))
-
+            q = self.wq_a(x)
+            q = self.wq_b(self.q_norm(q))
         # Use -1 instead of `n_heads` (or `n_kv_heads`) to infer the actual
         # local heads from sizes of q and kv as TP may have sharded them after
         # the above linear ops.
@@ -235,6 +241,24 @@ def forward(
         output = output.view(bsz, seqlen, -1)  # (bsz, seqlen, n_heads * v_head_dim)
         return self.wo(output)  # (bsz, seqlen, dim)
 
+    def init_weights(self, init_std: float):
+        linear_list = [
+            self.wkv_a,
+            self.wkv_b,
+        ]
+        if self.q_lora_rank > 0:
+            linear_list.extend([self.wq_a, self.wq_b])
+        else:
+            linear_list.append(self.wq)
+
+        for linear in linear_list:
+            nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02)
+        nn.init.trunc_normal_(self.wo.weight, mean=0.0, std=init_std)
+
+        self.kv_norm.reset_parameters()
+        if self.q_lora_rank > 0:
+            self.q_norm.reset_parameters()
+
 
 class FeedForward(nn.Module):
     """
@@ -266,7 +290,7 @@ def __init__(
     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):
+    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)
@@ -283,13 +307,16 @@ def __init__(self, layer_id: int, model_args: DeepSeekV3ModelArgs):
         self.attention = Attention(model_args)
         self.attention_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
         self.ffn_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
-        self.moe_enabled = layer_id < model_args.n_dense_layers
+        self.moe_enabled = layer_id >= model_args.n_dense_layers
 
         if self.moe_enabled:
             self.moe = MoE(model_args)
         else:
             self.feed_forward = FeedForward(model_args.dim, model_args.inter_dim)
 
+        # TODO: Need to revisit the weight initialization for the TransformerBlock
+        self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5
+
     def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):
         """
         Forward pass for the Transformer block.
@@ -308,6 +335,15 @@ def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):
             x = x + self.feed_forward(self.ffn_norm(x))
         return x
 
+    def init_weights(self, buffer_device: torch.device):
+        for norm in (self.attention_norm, self.ffn_norm):
+            norm.reset_parameters()
+        self.attention.init_weights(self.weight_init_std)
+        if self.moe_enabled:
+            self.moe.init_weights(self.weight_init_std, buffer_device)
+        else:
+            self.feed_forward.init_weights(self.weight_init_std)
+
 
 class DeepSeekV3Model(nn.Module, ModelProtocol):
     """
@@ -319,7 +355,7 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
         self.max_seq_len = model_args.max_seq_len
         self.tok_embeddings = nn.Embedding(model_args.vocab_size, model_args.dim)
         self.register_buffer(
-            "freqs_cis", precompute_freqs_cis(model_args), persistent=False
+            "freqs_cis", precompute_freqs_cis(model_args), persistent=True
         )
 
         self.layers = torch.nn.ModuleDict()
@@ -328,10 +364,36 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
 
         self.norm = nn.RMSNorm(model_args.dim)
         self.output = nn.Linear(
-            model_args.dim, model_args.vocab_size, dtype=torch.get_default_dtype()
+            model_args.dim,
+            model_args.vocab_size,
+            dtype=torch.get_default_dtype(),
+            bias=False,
         )
+        self.model_args = model_args
         self.init_weights()
 
+    def init_weights(self, buffer_device: torch.device | None = None) -> None:
+        buffer_device = buffer_device or self.freqs_cis.device
+        with torch.device(buffer_device):
+            self.freqs_cis = precompute_freqs_cis(self.model_args)
+        if self.tok_embeddings is not None:
+            nn.init.normal_(self.tok_embeddings.weight)
+        for layer in self.layers.values():
+            if layer is not None:
+                layer.init_weights(buffer_device=buffer_device)
+        if self.norm is not None:
+            self.norm.reset_parameters()
+        final_out_std = self.model_args.dim**-0.5
+        cutoff_factor = 3
+        if self.output is not None:
+            nn.init.trunc_normal_(
+                self.output.weight,
+                mean=0.0,
+                std=final_out_std,
+                a=-cutoff_factor * final_out_std,
+                b=cutoff_factor * final_out_std,
+            )
+
     def forward(self, tokens: torch.Tensor):
         """
         Forward pass for the Transformer model.
@@ -347,8 +409,5 @@ def forward(self, tokens: torch.Tensor):
         for layer in self.layers.values():
             h = layer(h, self.freqs_cis)
         h = self.norm(h)
-        output = self.output(h)  # (batch_size, seq_len, dim)
+        output = self.output(h)
         return output
-
-    def init_weights(self, buffer_device: torch.device | None = None) -> None:
-        pass

torchtitan/models/deepseek_v3/model/moe.py

@@ -89,6 +89,11 @@ def forward(
 
         return out
 
+    def init_weights(self, init_std: float):
+        nn.init.trunc_normal_(self.w1, mean=0.0, std=0.02)
+        nn.init.trunc_normal_(self.w2, mean=0.0, std=init_std)
+        nn.init.trunc_normal_(self.w3, mean=0.0, std=init_std)
+
 
 class TokenChoiceTopKRouter(nn.Module):
     """This class implements token-choice routing. In token-choice top-K routing, each token is
@@ -173,6 +178,9 @@ def forward(
 
         return top_scores, token_indices_experts_sorted, num_local_tokens_per_expert
 
+    def init_weights(self, init_std: float):
+        nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std)
+
 
 class MoE(nn.Module):
     def __init__(self, model_args: DeepSeekV3ModelArgs):
@@ -231,7 +239,6 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
         if self.load_balance_coeff is not None and self.load_balance_coeff > 0:
             self.register_full_backward_hook(self._update_expert_bias)
 
-    # TODO: double check the bias update logic. It aligns with the paper.
     def _update_expert_bias(self, *_):
         expert_bias_delta = self.load_balance_coeff * torch.sign(
             self.tokens_per_expert.mean() - self.tokens_per_expert

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -50,7 +50,7 @@ dataset = "c4_test"  # 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 = 1
+tensor_parallel_degree = 2
 enable_async_tensor_parallel = false
 
 [checkpoint]

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -38,7 +38,7 @@ decay_type = "linear"
 lr_min = 0.0
 
 [training]
-local_batch_size = 32
+local_batch_size = 16
 seq_len = 2048
 max_norm = 1.0  # grad norm clipping
 steps = 10
@@ -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 = 1
+tensor_parallel_degree = 2
 enable_async_tensor_parallel = false
 
 [checkpoint]