enable DSV3 manual bucketing

Author: ruisizhang123

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -38,9 +38,11 @@ def convert_modules_to_fqns(modules, module_to_fqn_mapping):
         result = []
         for m in modules:
             if isinstance(m, list):
-                result.append(convert_modules_to_fqns(m, module_to_fqn_mapping))
+                if fqn_list := convert_modules_to_fqns(m, module_to_fqn_mapping):
+                    result.append(fqn_list)
             else:
-                result.append(module_to_fqn_mapping.get(m, None))
+                if fqn := module_to_fqn_mapping.get(m):
+                    result.append(fqn)
         return result
 
     module_to_name = {m: n for n, m in model.named_modules()}

torchtitan/models/deepseek_v3/__init__.py

@@ -76,6 +76,31 @@
         attn_mask_type="block_causal",
     ),
     "16B": DeepSeekV3ModelArgs(
+        vocab_size=102400,
+        dim=2048,
+        inter_dim=10944,
+        moe_inter_dim=1408,
+        n_layers=27,
+        n_dense_layers=1,
+        n_heads=16,
+        moe_args=MoEArgs(
+            num_experts=64,
+            num_shared_experts=2,
+            top_k=6,
+            score_func="softmax",
+            route_norm=False,
+            score_before_experts=False,
+        ),
+        q_lora_rank=0,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        mscale=0.70,
+        use_flex_attn=False,
+        attn_mask_type="block_causal",
+    ),
+    "16B_flex_attn": DeepSeekV3ModelArgs(
         vocab_size=102400,
         dim=2048,
         inter_dim=10944,