clean

Author: wwwjn

torchtitan/models/deepseek_v3/model/model.py

@@ -368,15 +368,10 @@ def forward(self, tokens: torch.Tensor):
         Returns:
             torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
         """
-        print("Input tokens:", tokens)
         h = self.tok_embeddings(tokens)
-        print("After token embedding:", h)
 
         for layer in self.layers.values():
             h = layer(h, self.freqs_cis)
-            print(f"After layer {layer}: ", h)
         h = self.norm(h)
-        print("After normalization:", h)
         output = self.output(h)
-        print("Output logits:", output)
         return output

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 = 1
+steps = 10
 compile = false
 dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)