deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-04-19 10:08:59 -04:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Merge abaadd9b3e into 4cc6253d5c

Browse Source
This commit is contained in:
Saro 2025-04-09 11:49:55 +08:00 committed by GitHub
commit 3b10d4539f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
inference/model.py
View File

@ -115,16 +115,23 @@ class ParallelEmbedding(nn.Module):
Raises: Raises:
ValueError: If `world_size` is not defined. ValueError: If `world_size` is not defined.
""" """
if world_size > 1: def forward(self, x: torch.Tensor) -> torch.Tensor:
mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) if self.world_size < 1:
x = x - self.vocab_start_idx raise ValueError("world_size must be >= 1")
x[mask] = 0
y = F.embedding(x, self.weight)
if world_size > 1:
y[mask] = 0
dist.all_reduce(y)
return y
if self.world_size == 1:
return F.embedding(x, self.weight)
# For world_size > 1
mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
x = x - self.vocab_start_idx
x[mask] = 0
y = F.embedding(x, self.weight)
y[mask] = 0
dist.all_reduce(y)
return y
def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor: def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
""" """