[Qwen3] Fix model arg computation for MoE model (#1704)

as titled

Author: wwwjn

torchtitan/components/metrics.py

@@ -163,10 +163,12 @@ def close(self) -> None:
         if self.wandb.run is not None:
             self.wandb.finish()
 
+
 class LoggerContainer(BaseLogger):
     """Container to call all loggers enabled in the job config."""
+
     def __init__(self) -> None:
-        self._loggers : list[BaseLogger] = []
+        self._loggers: list[BaseLogger] = []
 
     def add_logger(self, logger_instance: BaseLogger) -> None:
         self._loggers.append(logger_instance)
@@ -183,6 +185,7 @@ def close(self) -> None:
         for logger_instance in self._loggers:
             logger_instance.close()
 
+
 def ensure_pp_loss_visible(
     parallel_dims: ParallelDims, job_config: JobConfig, color: Color
 ) -> None:

torchtitan/experiments/qwen3/model/args.py

@@ -55,11 +55,36 @@ def update_from_config(self, job_config: JobConfig, **kwargs) -> None:
         self.max_seq_len = seq_len
 
     def get_nparams_and_flops(self, model: nn.Module, seq_len: int) -> tuple[int, int]:
-        nparams = sum(p.numel() for p in model.parameters())
-        nparams_embedding = sum(
-            sum(p.numel() for p in m.parameters())
-            for m in model.children()
-            if isinstance(m, nn.Embedding)
+        nparams_embedding = 0
+        nparams_moe_router = 0
+        nparams_shared_experts = 0
+        nparams_experts = 0
+        nparams_dense = 0
+
+        for name, p in model.named_parameters():
+            if "embedding" in name:
+                nparams_embedding += p.numel()
+                nparams_dense += p.numel()
+            elif "moe.shared_experts" in name:
+                nparams_shared_experts += p.numel()
+            elif "moe.router" in name:
+                nparams_moe_router += p.numel()
+            elif "moe.experts" in name:
+                nparams_experts += p.numel()
+            else:
+                nparams_dense += p.numel()
+
+        nparams_sparse = nparams_moe_router + nparams_shared_experts + nparams_experts
+        nparams = nparams_dense + nparams_sparse
+        nparams_sparse_active = (
+            nparams_moe_router
+            + nparams_shared_experts
+            + nparams_experts * self.moe_args.top_k // self.moe_args.num_experts
+        )
+
+        logger.info(
+            f"Total parameter count: dense {nparams_dense:,}, "
+            f"sparse {nparams_sparse:,}, active {nparams_dense + nparams_sparse_active:,}"
         )
 
         l, h, q, t = (
@@ -68,10 +93,18 @@ def get_nparams_and_flops(self, model: nn.Module, seq_len: int) -> tuple[int, in
             self.dim // self.n_heads,
             seq_len,
         )
-        num_flops_per_token = 6 * (nparams - nparams_embedding) + 12 * l * h * q * t
+        # Reasoning behind the factor of 12 for the self-attention part of the formula:
+        # 1. each self-attention has 2 matmul in the forward and 4 in the backward (6)
+        # 2. the flash attention does 1 more matmul recomputation in the backward
+        #    but recomputation should not be counted in calculating MFU           (+0)
+        # 3. each matmul performs 1 multiplication and 1 addition                 (*2)
+        # 4. we follow the convention and do not account for sparsity in causal attention
+        num_flops_per_token = (
+            6 * (nparams_dense - nparams_embedding + nparams_sparse_active)
+            + 12 * l * h * q * t
+        )
 
         if self.enable_weight_tying:
-            # exclude model.token_embedding parameters from nparams
             nparams = nparams - nparams_embedding
 
         return nparams, num_flops_per_token

torchtitan/experiments/qwen3/model/model.py

@@ -132,6 +132,7 @@ def __init__(self, model_args: Qwen3ModelArgs):
         )
         self.n_rep = self.n_heads // self.n_kv_heads
         self.head_dim = model_args.head_dim
+        self.scaling = self.head_dim**-0.5
 
         # RMSNorm added here to the here to include the q-k norm
         # This is one of the main differences between Llama3 and Qwen3
@@ -209,7 +210,7 @@ def forward(
         xk = keys.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
         xv = values.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
 
-        output = self.sdpa(xq, xk, xv)
+        output = self.sdpa(xq, xk, xv, scale=self.scaling)
 
         output = output.transpose(
             1, 2