[DSV3] Apply TP on DSV3 (#1341)

Mostly adapted from llama4, change the TP plan based on the difference
between deepseek-v3 and llama.

Thanks @tianyu-l for the detailed walk through about deepseek-v3
attention model and TP plan! This diff is currently based on #1324 , and
we want to extract the MoE model in DSV3 and llama4 in a shared place.

Now we have:
1. FSDP
2. Activation Checkpointing
3. TP
4. CP in progress (hang due to some reason)

1. Make CP work

There are minor issue with the numerical verification: With
deterministic seed, the loss is not identical. I used `AdamW` optimizer.

1. FSDP degree=4 (blue line)
2. FSDP degree=4, TP degree = 2 (orange line)

<img width="1368" alt="Screenshot 2025-07-01 at 5 38 50 PM"
src="https://github.com/user-attachments/assets/38d96d75-6868-4482-a603-b9e10c692ed9"
/>

With `Adam` optimizer, the loss is **exactly the same**:
<img width="1368" alt="Screenshot 2025-07-02 at 1 26 32 PM"
src="https://github.com/user-attachments/assets/6b501d3c-4841-42b1-95fd-3971b16a5eeb"
/>

---------

Co-authored-by: Tianyu Liu <lty@fb.com>

Author: wwwjn

torchtitan/components/tokenizer.py

@@ -11,6 +11,7 @@
 from typing import Any, Optional
 
 from tokenizers import AddedToken, Tokenizer as HfTokenizer
+
 from typing_extensions import override
 
 

torchtitan/experiments/llama4/train_configs/debug_model.toml

@@ -26,7 +26,8 @@ tokenizer_path = "./tests/assets/test_tiktoken.model"
 # converters = ["float8"]
 
 [optimizer]
-name = "AdamW"
+# TODO: AdamW has numerical issues when TP is used, need to fix it
+name = "Adam"
 lr = 4e-3
 eps = 1e-15
 

torchtitan/models/deepseek_v3/README.md

@@ -0,0 +1,6 @@
+Download tokenizer:
+
+```
+# DeepSeek tokenizer (automatically downloads tokenizer.json and tokenizer_config.json)
+python scripts/download_tokenizer.py --repo_id deepseek-ai/DeepSeek-V3
+```

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.components.tokenizer import build_hf_tokenizer
 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,10 +117,10 @@
         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,
+        build_tokenizer_fn=build_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
     )
 )

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -6,9 +6,19 @@
 
 import torch.nn as nn
 from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import Replicate, Shard
+from torch.distributed.tensor.parallel import (
+    ColwiseParallel,
+    parallelize_module,
+    PrepareModuleInput,
+    RowwiseParallel,
+    SequenceParallel,
+)
 
 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_ep_tp
 from torchtitan.models.llama3.infra.parallelize import apply_ac, apply_fsdp
 from torchtitan.tools.logging import logger
 
@@ -19,6 +29,47 @@ def parallelize_deepseekv3(
     parallel_dims: ParallelDims,
     job_config: JobConfig,
 ):
+
+    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. \
+                torch.compile is not supported yet, which is required for async TP."
+            )
+
+        enable_float8_linear = "float8" in job_config.model.converters
+        float8_is_rowwise = job_config.float8.recipe_name in (
+            "rowwise",
+            "rowwise_with_gw_hp",
+        )
+
+        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"
+            )
+
+        apply_tp(
+            model,
+            world_mesh["tp"],
+            loss_parallel=parallel_dims.loss_parallel_enabled,
+            enable_float8_tensorwise_tp=False,
+            enable_async_tp=False,
+        )
+
+        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)
 
@@ -48,3 +99,79 @@ def parallelize_deepseekv3(
             logger.info("Applied FSDP to the model")
 
     return model
+
+
+def apply_tp(
+    model: nn.Module,
+    tp_mesh: DeviceMesh,
+    loss_parallel: bool,
+    enable_float8_tensorwise_tp: bool,
+    enable_async_tp: bool,
+):
+    """Apply tensor parallelism."""
+    # 1. Parallelize the embedding and shard its outputs (which are the first
+    # transformer block's inputs)
+    # 2. Parallelize the root norm layer over the sequence dim
+    # 3. Parallelize the final linear output layer
+    parallelize_module(
+        model,
+        tp_mesh,
+        {
+            "tok_embeddings": RowwiseParallel(
+                input_layouts=Replicate(),
+                output_layouts=Shard(1),
+            ),
+            "norm": SequenceParallel(),
+            "output": ColwiseParallel(
+                input_layouts=Shard(1),
+                output_layouts=Shard(-1) if loss_parallel else Replicate(),
+                use_local_output=not loss_parallel,
+            ),
+        },
+    )
+
+    rowwise_parallel, colwise_parallel, prepare_module_input = (
+        RowwiseParallel,
+        ColwiseParallel,
+        PrepareModuleInput,
+    )
+
+    # Apply tensor + sequence parallelism to every transformer block
+    # NOTE: At the cost of model code change, we can accelerate Sequence Parallel
+    #       by folding (and unfolding) the batch dimension and the sequence dimension.
+    #       Examples can be found at https://github.com/pytorch/torchtitan/pull/437
+    for transformer_block in model.layers.values():
+        layer_plan = {
+            "attention_norm": SequenceParallel(),
+            "attention": prepare_module_input(
+                input_layouts=(Shard(1), None),
+                desired_input_layouts=(Replicate(), None),
+            ),
+            "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(),
+        }
+
+        parallelize_module(
+            module=transformer_block,
+            device_mesh=tp_mesh,
+            parallelize_plan=layer_plan,
+        )
+
+    logger.info(
+        f"Applied {'Float8 tensorwise ' if enable_float8_tensorwise_tp else ''}{'Async ' if enable_async_tp else ''}"
+        "Tensor Parallelism to the model"
+    )

torchtitan/models/deepseek_v3/model/args.py

@@ -75,8 +75,8 @@ class DeepSeekV3ModelArgs(BaseModelArgs):
     n_limited_groups: int = 1
     score_func: Literal["softmax", "sigmoid"] = "softmax"
     route_scale: float = 1.0
-    use_grouped_mm: bool = False
-    load_balance_coeff: float | None = 1e-3
+    use_grouped_mm: bool = True
+    load_balance_coeff: float = 1e-3
     # Multi-Head Latent Attention (MLA)
     q_lora_rank: int = 0
     kv_lora_rank: int = 512
@@ -97,7 +97,7 @@ def update_from_config(self, job_config: JobConfig, tokenizer: Tokenizer) -> Non
         """
         Update the model_config config from the given job config.
         """
-        self.vocab_size = tokenizer.n_words
+        self.vocab_size = tokenizer.vocab_size
         self.max_seq_len = job_config.training.seq_len
 
     def get_nparams_and_flops(self, model: nn.Module, seq_len: int) -> tuple[int, int]: