Skip to content

Commit b59f468

Browse files
sdmyzlpsdmyzlp
committed
Add DSv32
1 parent 0e88661 commit b59f468

6 files changed

Lines changed: 622 additions & 8 deletions

File tree

‎torchtitan/models/deepseek_v3/__init__.py‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ def _make_dsv3_attn_config(
157157
)
158158

159159

160-
def _build_dsv3_layers(
160+
def _build_dsv3_layers_impl(
161161
*,
162162
n_layers: int,
163163
n_dense_layers: int,
@@ -183,18 +183,17 @@ def _build_dsv3_layers(
183183
moe_comm_backend: str,
184184
non_blocking_capacity_factor: float | None,
185185
rope: RoPE.Config,
186+
attn_fn,
187+
**attn_fn_kwargs,
186188
) -> list[TransformerBlock.Config]:
187-
"""Build the list of per-layer TransformerBlock configs.
189+
"""Generic per-layer TransformerBlock config builder.
188190
189-
Layers with layer_id < n_dense_layers get a dense FeedForward and no MoE.
190-
Layers with layer_id >= n_dense_layers get a MoE and no FeedForward.
191-
192-
Router and expert inits are constructed per-layer so depth-scaled
193-
initializers are correct for each layer's position.
191+
``attn_fn(layer_id=layer_id, dim=dim, n_heads=n_heads, q_lora_rank=..., ...,
192+
rope=rope, **attn_fn_kwargs)`` produces the attention config.
194193
"""
195194
layers = []
196195
for layer_id in range(n_layers):
197-
attn_cfg = _make_dsv3_attn_config(
196+
attn_cfg = attn_fn(
198197
layer_id=layer_id,
199198
dim=dim,
200199
n_heads=n_heads,
@@ -206,6 +205,7 @@ def _build_dsv3_layers(
206205
mscale=mscale,
207206
attn_backend=attn_backend,
208207
rope=rope,
208+
**attn_fn_kwargs,
209209
)
210210

211211
if layer_id < n_dense_layers:
@@ -262,6 +262,14 @@ def _build_dsv3_layers(
262262
return layers
263263

264264

265+
def _build_dsv3_layers(**kwargs):
266+
"""V3-specific thin wrapper around ``_build_dsv3_layers_impl``."""
267+
return _build_dsv3_layers_impl(
268+
attn_fn=_make_dsv3_attn_config,
269+
**kwargs,
270+
)
271+
272+
265273
def _debugmodel(
266274
attn_backend: str,
267275
moe_comm_backend: str,

‎torchtitan/models/deepseek_v3_2/__init__.py‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
import dataclasses
8+
from functools import partial
9+
10+
import torch.nn as nn
11+
12+
from torchtitan.components.optimizer import register_moe_load_balancing_hook
13+
from torchtitan.distributed.pipeline_parallel import pipeline_llm
14+
from torchtitan.models.common import (
15+
ComplexRoPE,
16+
CosSinRoPE,
17+
LayerNorm,
18+
Linear,
19+
RMSNorm,
20+
)
21+
from torchtitan.protocols.model_spec import ModelSpec
22+
23+
from torchtitan.models.deepseek_v3.__init__ import (
24+
_build_dsv3_layers_impl,
25+
_EMBEDDING_INIT,
26+
_LINEAR_INIT,
27+
_NORM_INIT,
28+
_output_linear_init,
29+
_make_dsv3_attn_config,
30+
)
31+
32+
from .model import (
33+
Attention,
34+
DeepSeekV32Model,
35+
DSAFlexAttention,
36+
Indexer,
37+
)
38+
from .state_dict_adapter import DeepSeekV32StateDictAdapter
39+
40+
# V3.2 parallelization delegates entirely to V3.
41+
from torchtitan.models.deepseek_v3.parallelize import (
42+
parallelize_deepseekv3 as parallelize_deepseekv32,
43+
)
44+
45+
__all__ = ["parallelize_deepseekv32", "DeepSeekV32Model", "deepseekv32_configs"]
46+
47+
48+
def _make_dsv32_attn_config(
49+
layer_id: int, *, index_n_heads: int, index_head_dim: int, index_topk: int,
50+
**kwargs,
51+
) -> Attention.Config:
52+
"""V3 MLA + Indexer + DSAFlexAttention inner."""
53+
c = _make_dsv3_attn_config(layer_id=layer_id, **kwargs)
54+
rhd = c.qk_rope_head_dim
55+
indexer_cfg = Indexer.Config(
56+
dim=c.dim, q_lora_rank=c.q_lora_rank,
57+
index_n_heads=index_n_heads, index_head_dim=index_head_dim,
58+
rope_head_dim=rhd, index_topk=index_topk,
59+
wq_b=Linear.Config(
60+
in_features=c.q_lora_rank,
61+
out_features=index_n_heads * index_head_dim,
62+
param_init=_LINEAR_INIT,
63+
),
64+
wk=Linear.Config(
65+
in_features=c.dim, out_features=index_head_dim,
66+
param_init=_LINEAR_INIT,
67+
),
68+
k_norm=LayerNorm.Config(normalized_shape=index_head_dim),
69+
weights_proj=Linear.Config(
70+
in_features=c.dim, out_features=index_n_heads,
71+
param_init={"weight": partial(nn.init.normal_, std=1.0)},
72+
),
73+
# Reference applies non-interleaved (rotate-half) rope with the
74+
# same yarn-scaled freqs as MLA.
75+
rope=CosSinRoPE.Config(
76+
dim=rhd, max_seq_len=c.rope.max_seq_len, theta=c.rope.theta,
77+
scaling="yarn", rope_factor=c.rope.rope_factor,
78+
beta_fast=c.rope.beta_fast, beta_slow=c.rope.beta_slow,
79+
original_seq_len=c.rope.original_seq_len,
80+
),
81+
)
82+
return Attention.Config(
83+
dim=c.dim, n_heads=c.n_heads,
84+
q_lora_rank=c.q_lora_rank, kv_lora_rank=c.kv_lora_rank,
85+
qk_nope_head_dim=c.qk_nope_head_dim, qk_rope_head_dim=c.qk_rope_head_dim,
86+
v_head_dim=c.v_head_dim, mscale=c.mscale,
87+
wq=c.wq, wq_a=c.wq_a, wq_b=c.wq_b, q_norm=c.q_norm,
88+
wkv_a=c.wkv_a, kv_norm=c.kv_norm, wkv_b=c.wkv_b, wo=c.wo,
89+
inner_attention=DSAFlexAttention.Config(index_topk=index_topk),
90+
rope=dataclasses.replace(c.rope),
91+
indexer=indexer_cfg,
92+
)
93+
94+
95+
def _build_dsv32_layers(**kwargs):
96+
"""V32-specific thin wrapper around ``_build_dsv3_layers_impl``."""
97+
return _build_dsv3_layers_impl(
98+
attn_fn=_make_dsv32_attn_config,
99+
**kwargs,
100+
)
101+
102+
103+
# ---------------------------------------------------------------------------
104+
# Model flavors
105+
# ---------------------------------------------------------------------------
106+
107+
def _debugmodel(
108+
attn_backend: str, moe_comm_backend: str,
109+
non_blocking_capacity_factor: float | None = None,
110+
) -> DeepSeekV32Model.Config:
111+
dim = 256
112+
layers = _build_dsv32_layers(
113+
n_layers=6, n_dense_layers=1, dim=dim, n_heads=16,
114+
q_lora_rank=64, kv_lora_rank=128, qk_nope_head_dim=32,
115+
qk_rope_head_dim=16, v_head_dim=32, mscale=0.70,
116+
dense_hidden_dim=512, moe_hidden_dim=128, num_experts=8,
117+
num_shared_experts=2, router_top_k=3, router_score_func="softmax",
118+
attn_backend=attn_backend, moe_comm_backend=moe_comm_backend,
119+
non_blocking_capacity_factor=non_blocking_capacity_factor,
120+
rope=ComplexRoPE.Config(dim=16, max_seq_len=16384, theta=10000.0,
121+
scaling="yarn", rope_factor=40.0,
122+
beta_fast=32.0, beta_slow=1.0, original_seq_len=4096),
123+
index_n_heads=4, index_head_dim=32, index_topk=16,
124+
)
125+
from torchtitan.models.common import Embedding
126+
return DeepSeekV32Model.Config(
127+
vocab_size=2048, dim=dim,
128+
tok_embeddings=Embedding.Config(num_embeddings=2048, embedding_dim=dim,
129+
param_init=_EMBEDDING_INIT),
130+
norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
131+
lm_head=Linear.Config(in_features=dim, out_features=2048,
132+
param_init=_output_linear_init(dim)),
133+
layers=layers,
134+
)
135+
136+
137+
deepseekv32_configs = {
138+
"debugmodel": _debugmodel,
139+
}
140+
141+
142+
def model_registry(
143+
flavor: str, attn_backend: str = "flex",
144+
moe_comm_backend: str = "standard",
145+
non_blocking_capacity_factor: float | None = None,
146+
) -> ModelSpec:
147+
config = deepseekv32_configs[flavor](
148+
attn_backend=attn_backend, moe_comm_backend=moe_comm_backend,
149+
non_blocking_capacity_factor=non_blocking_capacity_factor,
150+
)
151+
return ModelSpec(
152+
name="deepseek_v3_2", flavor=flavor, model=config,
153+
parallelize_fn=parallelize_deepseekv32, pipelining_fn=pipeline_llm,
154+
post_optimizer_build_fn=register_moe_load_balancing_hook,
155+
state_dict_adapter=DeepSeekV32StateDictAdapter,
156+
)

‎torchtitan/models/deepseek_v3_2/config_registry.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from torchtitan.config import ParallelismConfig
8+
from torchtitan.models.deepseek_v3.config_registry import deepseek_v3_debugmodel
9+
10+
from . import model_registry
11+
12+
13+
def deepseek_v3_2_debugmodel_full_dtensor():
14+
config = deepseek_v3_debugmodel()
15+
config.model_spec = model_registry("debugmodel")
16+
config.parallelism = ParallelismConfig(
17+
tensor_parallel_degree=2,
18+
spmd_backend="full_dtensor",
19+
enable_sequence_parallel=True,
20+
)
21+
return config

0 commit comments

Comments
 (0)