refactor(torchtitan): rollback Titan to 99c0cb2(20250907) and stabilize trainer UTs (#262)

Co-authored-by: zhenhuang12 <[email protected]>
Co-authored-by: xiaobochen-amd <[email protected]>

Author: 3 people

.github/workflows/docker/Dockerfile

@@ -1,6 +1,6 @@
 # Base image
 # FROM docker.io/rocm/megatron-lm:v25.9_gfx942
-FROM docker.io/rocm/pyt-megatron-lm-jax-nightly-private:pytorch_rocm7.0_20251024
+FROM docker.io/rocm/primus:v25.9_gfx942
 
 # Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .)
 ARG PRIMUS_TURBO_COMMIT

examples/torchtitan/configs/MI300X/deepseek_v3_16b-pretrain.yaml

@@ -71,6 +71,10 @@ modules:
         enable: true
         components: ["loss"]     # ["model", "loss"]
 
+      primus_turbo:
+        enable_primus_turbo: true
+        enable_attention_float8: false
+
       # quantize:
       #   linear:
       #     float8:

examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml

@@ -52,7 +52,7 @@ modules:
         enable_async_tensor_parallel: false
         pipeline_parallel_degree: 1
         pipeline_parallel_schedule: "Interleaved1F1B"
-        expert_parallel_degree: 1
+        expert_parallel_degree: 8
         expert_tensor_parallel_degree: 1
 
       checkpoint:

primus/backends/torchtitan/models/deepseek_v3/__init__.py

@@ -0,0 +1,63 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import torch
+from torchtitan.models.deepseek_v3.model.model import Attention as TTAttention
+from torchtitan.models.deepseek_v3.model.model import apply_rotary_emb
+
+
+class Attention(TTAttention):
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor,
+    ):
+        """
+        Forward pass for the Multi-Head Latent Attention (MLA) Layer.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
+            freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
+
+        Returns:
+            torch.Tensor: Output tensor with the same shape as the input.
+        """
+        bsz, seqlen, _ = x.size()
+
+        # Query projection
+        if self.q_lora_rank == 0:
+            q = self.wq(x)  # (bsz, seqlen, n_heads * qk_head_dim)
+        else:
+            q = self.wq_a(x)
+            q = self.wq_b(self.q_norm(q))
+        # Use -1 instead of `n_heads` (or `n_kv_heads`) to infer the actual
+        # local heads from sizes of q and kv as TP may have sharded them after
+        # the above linear ops.
+        q = q.view(bsz, seqlen, -1, self.qk_head_dim)
+        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        q_pe = apply_rotary_emb(q_pe, freqs_cis)
+        q = torch.cat([q_nope, q_pe], dim=-1)  # (bsz, seqlen, n_heads, qk_head_dim)
+
+        # Key-value projection
+        kv = self.wkv_a(x)  # (bsz, seqlen, kv_lora_rank + qk_rope_head_dim)
+        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+
+        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)  # (bsz, seqlen, 1, qk_rope_head_dim)
+
+        kv = self.wkv_b(self.kv_norm(kv))  # (bsz, seqlen, n_heads * (qk_nope_head_dim + v_head_dim))
+        kv = kv.view(bsz, seqlen, -1, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        k = torch.cat(
+            [k_nope, k_pe.expand(-1, -1, self.n_heads, -1)], dim=-1
+        )  # (bsz, seqlen, n_heads, qk_head_dim)
+
+        q = q.view(bsz, seqlen, -1, self.qk_head_dim)
+        k = k.view(bsz, seqlen, -1, self.qk_head_dim)
+        v = v.view(bsz, seqlen, -1, self.v_head_dim)
+
+        output = self.sdpa(q, k, v)
+        output = output.view(bsz, seqlen, -1)
+        return self.wo(output)

primus/backends/torchtitan/models/deepseek_v3/model/__init__.py

@@ -5,20 +5,16 @@
 ###############################################################################
 
 import torch
-from torch.nn.attention.flex_attention import BlockMask
+
+# from torch.nn.attention.flex_attention import BlockMask
 from torchtitan.models.llama3.model.model import Attention as TTAttention
 from torchtitan.models.llama3.model.model import apply_rotary_emb
 
-AttentionMasksType = dict[str, BlockMask] | BlockMask
+# AttentionMasksType = dict[str, BlockMask] | BlockMask
 
 
 class Attention(TTAttention):
-    def forward(
-        self,
-        x: torch.Tensor,
-        freqs_cis: torch.Tensor,
-        attention_masks: AttentionMasksType | None,
-    ):
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):
         bs, seqlen, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
 
@@ -35,7 +31,8 @@ def forward(
         # xk = repeat_kv(xk, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
         # xv = repeat_kv(xv, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
 
-        output = self.inner_attention(xq, xk, xv)
+        # output = self.inner_attention(xq, xk, xv)
+        output = self.sdpa(xq, xk, xv)
 
         output = output.view(bs, seqlen, -1)
         return self.wo(output)

primus/backends/torchtitan/models/deepseek_v3/model/model.py

@@ -7,10 +7,7 @@
 import torch
 from torchtitan.config.job_config import JobConfig
 from torchtitan.distributed import ParallelDims
-from torchtitan.models.attention import (
-    FlexAttentionWrapper,
-    ScaledDotProductAttentionWrapper,
-)
+from torchtitan.models.attention import FlexAttention, ScaledDotProductAttention
 from torchtitan.protocols.model_converter import (
     ModelConverter,
     register_model_converter,
@@ -21,7 +18,7 @@ def replace_turbo_attention_modules(model: torch.nn.Module, backend_type: str, u
     from primus_turbo.pytorch.modules import TurboAttention  # TODO: import Check
 
     for name, module in model.named_children():
-        if isinstance(module, (FlexAttentionWrapper, ScaledDotProductAttentionWrapper)):
+        if isinstance(module, (FlexAttention, ScaledDotProductAttention)):
             setattr(
                 model,
                 name,

primus/backends/torchtitan/models/llama3/model/model.py

@@ -7,5 +7,4 @@ model:
   flavor: "70B"
   hf_assets_path: "meta-llama/Llama-3.1-8B"
   converters:
-    - quantize.linear.float8
-    - quantize.grouped_mm.float8
+    - "float8"

primus/backends/torchtitan/primus_turbo_extensions/primus_turbo_converter.py

@@ -7,5 +7,4 @@ model:
   flavor: "8B"
   hf_assets_path: "meta-llama/Llama-3.1-8B"
   converters:
-    - quantize.linear.float8
-    - quantize.grouped_mm.float8
+    - "float8"