Skip to content

Commit 811953e

Browse files
Matrix-Z97Matrix-Z97
committed
Update DeepSeek V4 configs and sharding
1 parent 7811710 commit 811953e

5 files changed

Lines changed: 368 additions & 19 deletions

File tree

‎torchtitan/models/deepseek_v4/__init__.py‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,8 +581,230 @@ def _debugmodel(
581581
)
582582

583583

584+
def _deepseek_v4_flash(
585+
moe_comm_backend: str = "standard",
586+
non_blocking_capacity_factor: float | None = None,
587+
enable_indexer_loss: bool = True,
588+
) -> DeepSeekV4Model.Config:
589+
dim = 4096
590+
n_layers = 43
591+
vocab_size = 129280
592+
n_heads = 64
593+
head_dim = 512
594+
rope_head_dim = 64
595+
q_lora_rank = 1024
596+
o_lora_rank = 1024
597+
n_groups = 8
598+
compress_ratios = (1, 1) + (4, 128) * 20 + (4,)
599+
window_size = 128
600+
norm_eps = 1e-6
601+
index_n_heads = 64
602+
index_head_dim = 128
603+
index_topk = 512
604+
moe_inter_dim = 2048
605+
num_experts = 256
606+
num_shared_experts = 1
607+
top_k = 6
608+
n_hash_layers = 3
609+
route_norm = True
610+
route_scale = 1.5
611+
load_balance_coeff = 1e-3
612+
hc_mult = 4
613+
sinkhorn_iters = 20
614+
hc_eps = 1e-6
615+
dense_layers = set()
616+
max_seq_len = 4096
617+
compress_rope_theta = 160000.0
618+
original_seq_len = 65536
619+
620+
rope = ComplexRoPE.Config(
621+
dim=rope_head_dim,
622+
max_seq_len=max_seq_len,
623+
theta=10000.0,
624+
scaling="none",
625+
)
626+
rope_compress = ComplexRoPE.Config(
627+
dim=rope_head_dim,
628+
max_seq_len=max_seq_len,
629+
theta=compress_rope_theta,
630+
scaling="yarn",
631+
rope_factor=16.0,
632+
beta_fast=32.0,
633+
beta_slow=1.0,
634+
original_seq_len=original_seq_len,
635+
)
636+
637+
layers = _build_v4_layers(
638+
n_layers=n_layers,
639+
dim=dim,
640+
n_heads=n_heads,
641+
head_dim=head_dim,
642+
rope_head_dim=rope_head_dim,
643+
q_lora_rank=q_lora_rank,
644+
o_lora_rank=o_lora_rank,
645+
n_groups=n_groups,
646+
compress_ratios=compress_ratios,
647+
window_size=window_size,
648+
norm_eps=norm_eps,
649+
index_n_heads=index_n_heads,
650+
index_head_dim=index_head_dim,
651+
index_topk=index_topk,
652+
moe_inter_dim=moe_inter_dim,
653+
num_experts=num_experts,
654+
num_shared_experts=num_shared_experts,
655+
top_k=top_k,
656+
vocab_size=vocab_size,
657+
n_hash_layers=n_hash_layers,
658+
route_norm=route_norm,
659+
route_scale=route_scale,
660+
load_balance_coeff=load_balance_coeff,
661+
moe_comm_backend=moe_comm_backend,
662+
non_blocking_capacity_factor=non_blocking_capacity_factor,
663+
rope=rope,
664+
rope_compress=rope_compress,
665+
hc_mult=hc_mult,
666+
sinkhorn_iters=sinkhorn_iters,
667+
hc_eps=hc_eps,
668+
dense_layers=dense_layers,
669+
enable_indexer_loss=enable_indexer_loss,
670+
)
671+
672+
return DeepSeekV4Model.Config(
673+
dim=dim,
674+
vocab_size=vocab_size,
675+
norm_eps=norm_eps,
676+
tok_embeddings=Embedding.Config(
677+
num_embeddings=vocab_size,
678+
embedding_dim=dim,
679+
param_init=_EMBEDDING_INIT,
680+
),
681+
norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
682+
lm_head=Linear.Config(
683+
in_features=dim,
684+
out_features=vocab_size,
685+
param_init=_output_linear_init(dim),
686+
),
687+
layers=layers,
688+
hc_mult=hc_mult,
689+
compress_ratios=compress_ratios,
690+
n_layers=n_layers,
691+
)
692+
693+
694+
def _deepseek_v4_pro(
695+
moe_comm_backend: str = "standard",
696+
non_blocking_capacity_factor: float | None = None,
697+
enable_indexer_loss: bool = True,
698+
) -> DeepSeekV4Model.Config:
699+
dim = 7168
700+
n_layers = 61
701+
vocab_size = 129280
702+
n_heads = 128
703+
head_dim = 512
704+
rope_head_dim = 64
705+
q_lora_rank = 1536
706+
o_lora_rank = 1024
707+
n_groups = 16
708+
compress_ratios = (128,) + (128, 4) * 30
709+
window_size = 128
710+
norm_eps = 1e-6
711+
index_n_heads = 64
712+
index_head_dim = 128
713+
index_topk = 1024
714+
moe_inter_dim = 3072
715+
num_experts = 384
716+
num_shared_experts = 1
717+
top_k = 6
718+
n_hash_layers = 3
719+
route_norm = True
720+
route_scale = 1.5
721+
load_balance_coeff = 1e-3
722+
hc_mult = 4
723+
sinkhorn_iters = 20
724+
hc_eps = 1e-6
725+
dense_layers = set()
726+
max_seq_len = 4096
727+
compress_rope_theta = 160000.0
728+
original_seq_len = 65536
729+
730+
rope = ComplexRoPE.Config(
731+
dim=rope_head_dim,
732+
max_seq_len=max_seq_len,
733+
theta=10000.0,
734+
scaling="none",
735+
)
736+
rope_compress = ComplexRoPE.Config(
737+
dim=rope_head_dim,
738+
max_seq_len=max_seq_len,
739+
theta=compress_rope_theta,
740+
scaling="yarn",
741+
rope_factor=16.0,
742+
beta_fast=32.0,
743+
beta_slow=1.0,
744+
original_seq_len=original_seq_len,
745+
)
746+
747+
layers = _build_v4_layers(
748+
n_layers=n_layers,
749+
dim=dim,
750+
n_heads=n_heads,
751+
head_dim=head_dim,
752+
rope_head_dim=rope_head_dim,
753+
q_lora_rank=q_lora_rank,
754+
o_lora_rank=o_lora_rank,
755+
n_groups=n_groups,
756+
compress_ratios=compress_ratios,
757+
window_size=window_size,
758+
norm_eps=norm_eps,
759+
index_n_heads=index_n_heads,
760+
index_head_dim=index_head_dim,
761+
index_topk=index_topk,
762+
moe_inter_dim=moe_inter_dim,
763+
num_experts=num_experts,
764+
num_shared_experts=num_shared_experts,
765+
top_k=top_k,
766+
vocab_size=vocab_size,
767+
n_hash_layers=n_hash_layers,
768+
route_norm=route_norm,
769+
route_scale=route_scale,
770+
load_balance_coeff=load_balance_coeff,
771+
moe_comm_backend=moe_comm_backend,
772+
non_blocking_capacity_factor=non_blocking_capacity_factor,
773+
rope=rope,
774+
rope_compress=rope_compress,
775+
hc_mult=hc_mult,
776+
sinkhorn_iters=sinkhorn_iters,
777+
hc_eps=hc_eps,
778+
dense_layers=dense_layers,
779+
enable_indexer_loss=enable_indexer_loss,
780+
)
781+
782+
return DeepSeekV4Model.Config(
783+
dim=dim,
784+
vocab_size=vocab_size,
785+
norm_eps=norm_eps,
786+
tok_embeddings=Embedding.Config(
787+
num_embeddings=vocab_size,
788+
embedding_dim=dim,
789+
param_init=_EMBEDDING_INIT,
790+
),
791+
norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
792+
lm_head=Linear.Config(
793+
in_features=dim,
794+
out_features=vocab_size,
795+
param_init=_output_linear_init(dim),
796+
),
797+
layers=layers,
798+
hc_mult=hc_mult,
799+
compress_ratios=compress_ratios,
800+
n_layers=n_layers,
801+
)
802+
803+
584804
deepseek_v4_configs = {
585805
"debugmodel": _debugmodel,
806+
"deepseek_v4_flash": _deepseek_v4_flash,
807+
"deepseek_v4_pro": _deepseek_v4_pro,
586808
}
587809

588810

‎torchtitan/models/deepseek_v4/attention.py‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@
55

66
from dataclasses import dataclass
77

8+
import spmd_types as spmd
89
import torch
910
import torch.nn.functional as F
1011
from torch import nn
1112

1213
from torchtitan.distributed.aux_loss import inject_aux_loss
13-
from torchtitan.distributed.utils import dist_max, dist_mean
14+
from torchtitan.distributed.utils import dist_max, dist_mean, get_spmd_backend
1415
from torchtitan.models.common.attention import (
1516
BaseAttention,
1617
create_attention_mask,
@@ -23,6 +24,14 @@
2324
from .compressor import Compressor, Indexer
2425

2526

27+
def _assert_spmd_attention_type(tensor, *, tp):
28+
if get_spmd_backend() == "spmd_types":
29+
spmd.assert_type(
30+
tensor,
31+
{"dp": spmd.S(0), "cp": spmd.S(1), "tp": tp},
32+
)
33+
34+
2635
class DSAIndexerAuxLoss(Module):
2736
@dataclass(kw_only=True, slots=True)
2837
class Config(Module.Config):
@@ -439,17 +448,24 @@ def forward(self, x, attention_masks=None, positions=None):
439448
rd = self.rope_head_dim
440449

441450
qr = self.q_norm(self.wq_a(x))
442-
q = self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim))
451+
_assert_spmd_attention_type(qr, tp=spmd.R)
452+
q = self.wq_b(qr)
453+
with spmd.local():
454+
q = q.view(bsz, seqlen, -1, self.head_dim)
455+
_assert_spmd_attention_type(q, tp=spmd.S(2))
443456
q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + self.norm_eps)
444457
q_nope, q_rope = torch.split(q, [self.head_dim - rd, rd], dim=-1)
445458

446459
kv = self.wkv(x)
447460
kv = self.kv_norm(kv)
461+
_assert_spmd_attention_type(kv, tp=spmd.R)
448462
kv_nope, kv_rope = torch.split(kv, [self.head_dim - rd, rd], dim=-1)
449463

450464
q_rope, kv_rope = self.rope(q_rope, kv_rope.unsqueeze(2), positions)
451465
q = torch.cat([q_nope, q_rope], dim=-1)
452466
kv = torch.cat([kv_nope, kv_rope.squeeze(2)], dim=-1)
467+
_assert_spmd_attention_type(q, tp=spmd.S(2))
468+
_assert_spmd_attention_type(kv, tp=spmd.R)
453469

454470
kv_compress = compress_topk_idxs = index_score = None
455471

@@ -477,10 +493,12 @@ def forward(self, x, attention_masks=None, positions=None):
477493
attn_sink_param = self.attn_sink.weight.squeeze(-1)
478494
if kv_compress is None:
479495
kv_compress = kv.new_empty((bsz, 0, self.head_dim))
496+
_assert_spmd_attention_type(kv_compress, tp=spmd.R)
480497
if compress_topk_idxs is None:
481498
compress_topk_idxs = torch.empty(
482499
(bsz, seqlen, 0), dtype=torch.int64, device=x.device
483500
)
501+
_assert_spmd_attention_type(compress_topk_idxs, tp=spmd.R)
484502
attn_out = self.inner_attention(
485503
q, kv, attn_sink_param, kv_compress, compress_topk_idxs,
486504
)
@@ -506,11 +524,17 @@ def forward(self, x, attention_masks=None, positions=None):
506524
o_nope, o_rope = torch.split(o, [self.head_dim - rd, rd], dim=-1)
507525
o_rope = self.rope(o_rope, o_rope, positions)[0]
508526
o = torch.cat([o_nope, o_rope], dim=-1)
509-
510-
n_local_groups = self.n_groups // (self.n_heads // o.shape[2])
511-
o = o.view(bsz, seqlen, n_local_groups, -1)
512-
# wo_a is a Linear module; access its weight directly for the grouped
513-
# einsum (not a standard Linear forward).
514-
wo_a = self.wo_a.weight.view(n_local_groups, self.o_lora_rank, -1)
527+
_assert_spmd_attention_type(o, tp=spmd.S(2))
528+
529+
with spmd.local():
530+
n_local_groups = self.n_groups // (self.n_heads // o.shape[2])
531+
o = o.view(bsz, seqlen, n_local_groups, -1)
532+
_assert_spmd_attention_type(o, tp=spmd.S(2))
533+
# wo_a is a Linear module; access its weight directly for the grouped
534+
# einsum (not a standard Linear forward).
535+
wo_a = self.wo_a.weight.view(n_local_groups, self.o_lora_rank, -1)
515536
o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
516-
return self.wo_b(o.reshape(bsz, seqlen, -1))
537+
with spmd.local():
538+
o = o.reshape(bsz, seqlen, -1)
539+
_assert_spmd_attention_type(o, tp=spmd.S(2))
540+
return self.wo_b(o)

‎torchtitan/models/deepseek_v4/compressor.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,25 @@
77
from dataclasses import dataclass
88
from functools import partial
99

10+
import spmd_types as spmd
1011
import torch
1112
import torch.nn.functional as F
1213
from torch import nn
1314

15+
from torchtitan.distributed.utils import get_spmd_backend
1416
from torchtitan.models.common.nn_modules import Linear, RMSNorm
1517
from torchtitan.models.common.rope import RoPE
1618
from torchtitan.protocols.module import Module
1719

1820

21+
def _assert_spmd_replicated_activation(tensor):
22+
if get_spmd_backend() == "spmd_types":
23+
spmd.assert_type(
24+
tensor,
25+
{"dp": spmd.S(0), "cp": spmd.S(1), "tp": spmd.R},
26+
)
27+
28+
1929
def _make_hadamard_mat(n: int, device: torch.device | str | None = None) -> torch.Tensor:
2030
n_pow2 = 2 ** math.ceil(math.log2(n))
2131
H = torch.tensor([[1.0, 1.0], [1.0, -1.0]], device=device)
@@ -96,6 +106,7 @@ def forward(self, x, positions=None):
96106
kv_rope.unsqueeze(2), kv_rope.unsqueeze(2), comp_positions
97107
)[0]
98108
kv = torch.cat([kv_nope, kv_rope.squeeze(2)], dim=-1)
109+
_assert_spmd_replicated_activation(kv)
99110
return kv
100111

101112

@@ -166,10 +177,13 @@ def forward(
166177
bsz, seqlen, _ = x.size()
167178
rd = self.rope_head_dim
168179
q = self.wq_b(qr)
169-
q = q.view(bsz, seqlen, self.num_index_heads, self.head_dim)
180+
with spmd.local():
181+
q = q.view(bsz, seqlen, self.num_index_heads, self.head_dim)
182+
_assert_spmd_replicated_activation(q)
170183
q_nope, q_rope = torch.split(q, [self.head_dim - rd, rd], dim=-1)
171184
q_rope = self.rope(q_rope, q_rope, positions)[0]
172185
q = torch.cat([q_nope, q_rope], dim=-1)
186+
_assert_spmd_replicated_activation(q)
173187
hadamard_mat = self.hadamard_mat.to(device=q.device, dtype=q.dtype)
174188
q = self._rotate_activation(q, hadamard_mat)
175189
k = self.compressor(x, positions=positions)
@@ -187,4 +201,6 @@ def forward(
187201
)
188202
mask = topk_idxs >= compress_causal_limit
189203
compress_topk_idxs = torch.where(mask, -1, topk_idxs + offset)
204+
_assert_spmd_replicated_activation(compress_topk_idxs)
205+
_assert_spmd_replicated_activation(index_score)
190206
return compress_topk_idxs, index_score

0 commit comments

Comments
 (0)