add TP v1

Author: wwwjn

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -32,8 +32,10 @@ def parallelize_deepseekv3(
 
     if parallel_dims.tp_enabled:
         if job_config.parallelism.enable_async_tensor_parallel:
+            # TODO(jianiw): This branch needs to be tested and enabled
             raise NotImplementedError(
-                "Currently, async TP is not tested for deepseekv3"
+                "Currently, async TP is not tested for deepseekv3. \
+                torch.compile is not supported yet, which is required for async TP."
             )
 
         enable_float8_linear = "float8" in job_config.model.converters
@@ -44,15 +46,11 @@ def parallelize_deepseekv3(
 
         enable_float8_tensorwise_tp = enable_float8_linear and not float8_is_rowwise
         if enable_float8_tensorwise_tp:
+            # TODO(jianiw): This branch needs to be tested and enabled
             raise NotImplementedError(
                 "Currently, float8 tensorwise TP is not tested for deepseekv3"
             )
 
-        if parallel_dims.loss_parallel_enabled:
-            raise NotImplementedError(
-                "Currently, loss parallel is not tested for deepseekv3"
-            )
-
         apply_tp(
             model,
             world_mesh["tp"],

torchtitan/models/deepseek_v3/model/model.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
-from re import I
 from typing import Tuple
 
 import torch

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -50,6 +50,8 @@ dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
+tensor_parallel_degree = 1
+enable_async_tensor_parallel = false
 
 [checkpoint]
 enable_checkpoint = false

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -49,8 +49,8 @@ dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 2
-disable_loss_parallel = true
+tensor_parallel_degree = 1
+enable_async_tensor_parallel = false
 
 [checkpoint]
 enable_checkpoint = false