add 671B model config

Author: wwwjn

torchtitan/models/deepseek_v3/README.md

@@ -27,6 +27,11 @@ CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ./r
 CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml" ./run_train.sh
 ```
 
+```bash
+# 671B parameter model
+CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/deepseek_v3_671b.toml" ./run_train.sh
+```
+
 
 ## Supported Features
 - FSDP, HSDP
@@ -36,6 +41,8 @@ CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml"
 
 
 ## To be added
+- TP:
+    - TP has a known numerical issue with DeepSeek-V3 (https://github.com/pytorch/torchtitan/pull/1373#issuecomment-3050249520).
 - Modeling
     - Merge DeepSeek-V3 and Llama4 MoE common components
 - Parallelism

torchtitan/models/deepseek_v3/__init__.py

@@ -30,8 +30,8 @@
     "debugmodel": DeepSeekV3ModelArgs(
         vocab_size=102400,
         dim=256,
-        inter_dim=10944,
-        moe_inter_dim=1408,
+        inter_dim=1024,
+        moe_inter_dim=256,
         n_layers=3,
         n_dense_layers=1,
         n_heads=16,

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -11,6 +11,7 @@
     ColwiseParallel,
     parallelize_module,
     PrepareModuleInput,
+    PrepareModuleInputOutput,
     RowwiseParallel,
     SequenceParallel,
 )
@@ -51,7 +52,7 @@ def parallelize_deepseekv3(
                 "Currently, float8 tensorwise TP is not tested for deepseekv3"
             )
 
-        apply_tp(
+        apply_non_moe_tp(
             model,
             world_mesh["tp"],
             loss_parallel=parallel_dims.loss_parallel_enabled,
@@ -133,7 +134,7 @@ def parallelize_deepseekv3(
     return model
 
 
-def apply_tp(
+def apply_non_moe_tp(
     model: nn.Module,
     tp_mesh: DeviceMesh,
     loss_parallel: bool,
@@ -145,6 +146,7 @@ def apply_tp(
     # transformer block's inputs)
     # 2. Parallelize the root norm layer over the sequence dim
     # 3. Parallelize the final linear output layer
+    logger.warning("There are known issue with TP for deepseekv3. Please see details in discussion: https://github.com/pytorch/torchtitan/pull/1373#issuecomment-3050249520.")
     parallelize_module(
         model,
         tp_mesh,
@@ -182,21 +184,36 @@ def apply_tp(
             "attention.wkv_a": NoParallel(),
             "attention.wkv_b": colwise_parallel(),
             "attention.kv_norm": NoParallel(),
-            "attention.wq_a": NoParallel(),
-            "attention.wq_b": colwise_parallel(),
-            "attention.q_norm": NoParallel(),
-            "attention.wq": colwise_parallel(),  # This is only used when q_lora_rank==0
             "attention.wo": rowwise_parallel(output_layouts=Shard(1)),
             "ffn_norm": SequenceParallel(),
-            "feed_forward": prepare_module_input(
-                input_layouts=(Shard(1),),
-                desired_input_layouts=(Replicate(),),
-            ),
-            "feed_forward.w1": colwise_parallel(),
-            "feed_forward.w2": rowwise_parallel(output_layouts=Shard(1)),
-            "feed_forward.w3": colwise_parallel(),
         }
 
+        if transformer_block.attention.q_lora_rank == 0:
+            layer_plan.update(
+                {
+                    "attention.wq": colwise_parallel(),  # This is only used when q_lora_rank==0
+                }
+            )
+        else:
+            layer_plan.update(
+                {
+                    "attention.wq_a": NoParallel(),
+                    "attention.wq_b": colwise_parallel(),
+                    "attention.q_norm": NoParallel(),
+                }
+            )
+
+        if not transformer_block.moe_enabled:
+            layer_plan.update({
+                "feed_forward": prepare_module_input(
+                    input_layouts=(Shard(1),),
+                    desired_input_layouts=(Replicate(),),
+                ),
+                "feed_forward.w1": colwise_parallel(),
+                "feed_forward.w2": rowwise_parallel(output_layouts=Shard(1)),
+                "feed_forward.w3": colwise_parallel(),
+            })
+
         parallelize_module(
             module=transformer_block,
             device_mesh=tp_mesh,

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -24,7 +24,7 @@ enable_wandb = false
 name = "deepseek_v3"
 flavor = "debugmodel"
 # test tokenizer.model, for debug purpose only
-tokenizer_path = "./assets/tokenizer/DeepSeek-V3"
+tokenizer_path = "./tests/assets/tokenizer"
 # converters = ["float8"]
 
 [optimizer]
@@ -40,7 +40,7 @@ lr_min = 0.0
 
 [training]
 local_batch_size = 8
-seq_len = 4096
+seq_len = 2048
 max_norm = 1.0  # grad norm clipping
 steps = 10
 compile = false
@@ -69,4 +69,5 @@ selective_ac_option = '2'  # 'int' = ac every positive int layer or 'op', ac bas
 [float8]
 enable_fsdp_float8_all_gather = false
 precompute_float8_dynamic_scale_for_fsdp = false
-filter_fqns = ["output"]
+filter_fqns = ["output", "router.gate"]
+moe_fqns = ["experts"]

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -22,26 +22,25 @@ enable_wandb = false
 [model]
 name = "deepseek_v3"
 flavor = "16B"
-# test tokenizer.model, for debug purpose only
 tokenizer_path = "./assets/tokenizer/DeepSeek-V3"
 # converters = ["float8"]
 
 [optimizer]
 name = "AdamW"
-lr = 8e-4
+lr = 2.2e-4
 eps = 1e-8
 
 [lr_scheduler]
-warmup_steps = 2  # lr scheduler warm up, normally 20% of the train steps
+warmup_steps = 200  # lr scheduler warm up, normally 20% of the train steps
 decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
 decay_type = "linear"
-lr_min = 0.0
+lr_min = 2.2e-5
 
 [training]
 local_batch_size = 8
 seq_len = 4096
 max_norm = 1.0  # grad norm clipping
-steps = 100
+steps = 1000
 compile = false
 dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
 
@@ -67,4 +66,5 @@ mode = "full"  # ["none", "selective", "full"]
 [float8]
 enable_fsdp_float8_all_gather = false
 precompute_float8_dynamic_scale_for_fsdp = false
-filter_fqns = ["output"]
+filter_fqns = ["output", "router.gate"]
+moe_fqns = ["experts"]

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_671b.toml

@@ -0,0 +1,70 @@
+# torchtitan Config.toml
+
+[job]
+dump_folder = "./outputs"
+description = "DeepSeek-V3 671B model training"
+print_args = false
+
+[profiling]
+enable_profiling = false
+save_traces_folder = "profile_trace"
+profile_freq = 10
+enable_memory_snapshot = false
+save_memory_snapshot_folder = "memory_snapshot"
+
+[metrics]
+log_freq = 10
+disable_color_printing = false
+enable_tensorboard = false
+save_tb_folder = "tb"
+enable_wandb = false
+
+[model]
+name = "deepseek_v3"
+flavor = "671B"
+tokenizer_path = "./assets/tokenizer/DeepSeek-V3"
+# converters = ["float8"]
+
+[optimizer]
+name = "AdamW"
+lr = 2.2e-4
+eps = 1e-8
+
+[lr_scheduler]
+warmup_steps = 2_000  # lr scheduler warm up, normally 20% of the train steps
+decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
+decay_type = "linear"
+lr_min = 2.2e-5
+
+[training]
+local_batch_size = 4
+seq_len = 4096
+max_norm = 1.0  # grad norm clipping
+steps = 10_000
+compile = false
+dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
+
+[parallelism]
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = -1
+fsdp_reshard_after_forward = "default" # default / never / always
+tensor_parallel_degree = 8
+enable_async_tensor_parallel = false
+expert_parallel_degree = 1
+
+[checkpoint]
+enable_checkpoint = false
+folder = "checkpoint"
+interval = 500
+last_save_model_weights_only = false
+export_dtype = "float32"
+async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem]"
+
+[activation_checkpoint]
+mode = "full"  # ["none", "selective", "full"]
+
+[float8]
+enable_fsdp_float8_all_gather = false
+precompute_float8_dynamic_scale_for_fsdp = false
+filter_fqns = ["output", "router.gate"]
+moe_fqns = ["experts"]