add dsv32

Author: sdmyzlp

torchtitan/models/deepseek_v3/__init__.py

@@ -15,7 +15,9 @@
 from torchtitan.distributed.pipeline_parallel import pipeline_llm
 from torchtitan.models.common import (
     ComplexRoPE,
+    CosSinRoPE,
     Embedding,
+    LayerNorm,
     Linear,
     RMSNorm,
     RoPE,
@@ -33,7 +35,13 @@
 from torchtitan.protocols.model import ModelConfigConverter
 from torchtitan.protocols.model_spec import ModelSpec
 
-from .model import Attention, DeepSeekV3Model, DeepSeekV3TransformerBlock
+from .model import (
+    Attention,
+    DeepSeekV3Model,
+    DeepSeekV3TransformerBlock,
+    DeepSeekSparseAttention,
+    Indexer,
+)
 from .parallelize import parallelize_deepseekv3
 from .state_dict_adapter import DeepSeekV3StateDictAdapter
 
@@ -88,14 +96,50 @@ def _make_dsv3_attn_config(
     mscale: float = 1.0,
     attn_backend: str,
     rope: RoPE.Config,
+    index_n_heads: int | None = None,
+    index_head_dim: int | None = None,
+    index_topk: int | None = None,
 ) -> Attention.Config:
     """Build a fully-specified DeepSeek V3 MLA Attention.Config.
 
     All Linear and RMSNorm sub-configs have their dimensional fields set.
     When q_lora_rank == 0, sets wq (not wq_a/wq_b).
     When q_lora_rank > 0, sets wq_a/wq_b (not wq).
+    When index_* kwargs are provided, also builds the Lightning Indexer
+    sub-modules and wraps the inner attention with DeepSeekSparseAttention.
     """
     inner_attention = get_attention_config(attn_backend)
+    indexer = None
+    if index_n_heads is not None:
+        assert index_head_dim is not None and index_topk is not None
+        indexer = Indexer.Config(
+            dim=dim, q_lora_rank=q_lora_rank,
+            index_n_heads=index_n_heads, index_head_dim=index_head_dim,
+            rope_head_dim=qk_rope_head_dim, index_topk=index_topk,
+            wq_b=Linear.Config(
+                in_features=q_lora_rank,
+                out_features=index_n_heads * index_head_dim,
+                param_init=_LINEAR_INIT,
+            ),
+            wk=Linear.Config(
+                in_features=dim, out_features=index_head_dim,
+                param_init=_LINEAR_INIT,
+            ),
+            k_norm=LayerNorm.Config(normalized_shape=index_head_dim),
+            weights_proj=Linear.Config(
+                in_features=dim, out_features=index_n_heads,
+                param_init={"weight": partial(nn.init.normal_, std=1.0)},
+            ),
+            rope=CosSinRoPE.Config(
+                dim=qk_rope_head_dim, max_seq_len=rope.max_seq_len,
+                theta=rope.theta, scaling=rope.scaling,
+                rope_factor=rope.rope_factor,
+                beta_fast=rope.beta_fast, beta_slow=rope.beta_slow,
+                original_seq_len=rope.original_seq_len,
+            ),
+        )
+        inner_attention = DeepSeekSparseAttention.Config(index_topk=index_topk)
+
     qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
 
     if q_lora_rank == 0:
@@ -154,6 +198,7 @@ def _make_dsv3_attn_config(
         ),
         inner_attention=inner_attention,
         rope=dataclasses.replace(rope),
+        indexer=indexer,
     )
 
 
@@ -183,6 +228,9 @@ def _build_dsv3_layers(
     moe_comm_backend: str,
     non_blocking_capacity_factor: float | None,
     rope: RoPE.Config,
+    index_n_heads: int | None = None,
+    index_head_dim: int | None = None,
+    index_topk: int | None = None,
 ) -> list[TransformerBlock.Config]:
     """Build the list of per-layer TransformerBlock configs.
 
@@ -206,6 +254,9 @@ def _build_dsv3_layers(
             mscale=mscale,
             attn_backend=attn_backend,
             rope=rope,
+            index_n_heads=index_n_heads,
+            index_head_dim=index_head_dim,
+            index_topk=index_topk,
         )
 
         if layer_id < n_dense_layers:
@@ -523,8 +574,83 @@ def _671b(
     )
 
 
+def _debugmodel_v32(
+    attn_backend: str,
+    moe_comm_backend: str,
+    non_blocking_capacity_factor: float | None = None,
+) -> DeepSeekV3Model.Config:
+    dim = 256
+    n_layers = 6
+    vocab_size = 2048
+    n_heads = 16
+    q_lora_rank = 64
+    kv_lora_rank = 128
+    qk_nope_head_dim = 64
+    qk_rope_head_dim = 64
+    v_head_dim = 64
+    index_n_heads = 4
+    index_head_dim = 128
+    index_topk = 32
+    moe_hidden_dim = 256
+    num_shared_experts = 2
+    dense_hidden_dim = 1024
+    num_experts = 8
+    n_dense_layers = 1
+
+    layers = _build_dsv3_layers(
+        n_layers=n_layers,
+        n_dense_layers=n_dense_layers,
+        dim=dim,
+        n_heads=n_heads,
+        q_lora_rank=q_lora_rank,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        mscale=0.70,
+        dense_hidden_dim=dense_hidden_dim,
+        moe_hidden_dim=moe_hidden_dim,
+        num_experts=num_experts,
+        num_shared_experts=num_shared_experts,
+        router_top_k=3,
+        router_score_func="softmax",
+        aux_loss_coeff=1e-4,
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+        rope=ComplexRoPE.Config(
+            dim=qk_rope_head_dim,
+            max_seq_len=4096 * 4,
+            theta=10000.0,
+            scaling="yarn",
+            rope_factor=40.0,
+            beta_fast=32.0,
+            beta_slow=1.0,
+            original_seq_len=4096,
+        ),
+        index_n_heads=index_n_heads,
+        index_head_dim=index_head_dim,
+        index_topk=index_topk,
+    )
+    return DeepSeekV3Model.Config(
+        vocab_size=vocab_size,
+        dim=dim,
+        tok_embeddings=Embedding.Config(
+            num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT
+        ),
+        norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
+        lm_head=Linear.Config(
+            in_features=dim,
+            out_features=vocab_size,
+            param_init=_output_linear_init(dim),
+        ),
+        layers=layers,
+    )
+
+
 deepseekv3_configs = {
     "debugmodel": _debugmodel,
+    "debugmodel_v32": _debugmodel_v32,
     "16B": _16b,
     "236B": _236b,
     "671B": _671b,

torchtitan/models/deepseek_v3/config_registry.py

@@ -73,6 +73,41 @@ def deepseek_v3_debugmodel_hybridep() -> Trainer.Config:
     return config
 
 
+def deepseek_v3_debugmodel_v32() -> Trainer.Config:
+    model_spec = model_registry("debugmodel_v32")
+    return Trainer.Config(
+        loss=ChunkedLossWrapper.Config(
+            loss_fn=CrossEntropyLoss.Config(
+                global_vocab_size=decoder_vocab_size(model_spec),
+            ),
+        ),
+        hf_assets_path="./tests/assets/tokenizer",
+        metrics=MetricsProcessor.Config(log_freq=1),
+        model_spec=model_spec,
+        dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"),
+        optimizer=default_adamw(lr=8e-4),
+        lr_scheduler=LRSchedulersContainer.Config(
+            warmup_steps=2,
+            decay_ratio=0.8,
+            decay_type="linear",
+            min_lr_factor=0.0,
+        ),
+        training=TrainingConfig(
+            local_batch_size=4,
+            seq_len=2048,
+            steps=10,
+        ),
+        parallelism=ParallelismConfig(
+            expert_parallel_degree=1,
+        ),
+        checkpoint=CheckpointManager.Config(
+            interval=10,
+            last_save_model_only=False,
+        ),
+        activation_checkpoint=SelectiveAC.Config(),
+    )
+
+
 def deepseek_v3_debugmodel_minimal_async_ep() -> Trainer.Config:
     config = deepseek_v3_debugmodel()
     config.model_spec = model_registry(