rebase onto #1324

Author: wwwjn

torchtitan/models/deepseek_v3/__init__.py

@@ -8,9 +8,9 @@
 
 from torchtitan.components.loss import build_cross_entropy_loss
 from torchtitan.components.lr_scheduler import build_lr_schedulers
-from torchtitan.components.optimizer import build_optimizers
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.datasets.tokenizer.tiktoken import build_tiktoken_tokenizer
+from torchtitan.experiments.llama4.optimizer import build_llama4_optimizers
 
 from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
 
@@ -117,7 +117,7 @@
         config=deepseekv3_configs,
         parallelize_fn=parallelize_deepseekv3,
         pipelining_fn=None,
-        build_optimizers_fn=build_optimizers,
+        build_optimizers_fn=build_llama4_optimizers,  # use optimizer hooks to update expert weights
         build_lr_schedulers_fn=build_lr_schedulers,
         build_dataloader_fn=build_hf_dataloader,
         build_tokenizer_fn=build_tiktoken_tokenizer,

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -18,7 +18,7 @@
 from torchtitan.config_manager import JobConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed import ParallelDims
 from torchtitan.experiments.llama4.infra.expert_parallel import NoParallel
-from torchtitan.experiments.llama4.infra.parallelize import apply_moe_tp
+from torchtitan.experiments.llama4.infra.parallelize import apply_moe_ep_tp
 from torchtitan.models.llama3.infra.parallelize import apply_ac, apply_fsdp
 from torchtitan.tools.logging import logger
 
@@ -59,7 +59,16 @@ def parallelize_deepseekv3(
             enable_async_tp=False,
         )
 
-        apply_moe_tp(model, world_mesh["tp"])
+        apply_moe_ep_tp(
+            model,
+            tp_mesh=world_mesh["tp"] if parallel_dims.tp_enabled else None,
+            ep_mesh=world_mesh["ep"] if parallel_dims.ep_enabled else None,
+            ep_tp_mesh=(
+                world_mesh["ep", "tp"]
+                if parallel_dims.tp_enabled and parallel_dims.ep_enabled
+                else None
+            ),
+        )
 
     if job_config.activation_checkpoint.mode != "none":
         apply_ac(model, job_config.activation_checkpoint)

torchtitan/models/deepseek_v3/model/args.py

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

torchtitan/models/deepseek_v3/model/model.py

@@ -8,7 +8,6 @@
 from typing import Tuple
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 from torchtitan.models.attention import build_attention
 from torchtitan.protocols.train_spec import ModelProtocol
@@ -369,10 +368,15 @@ 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