deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-05-23 02:36:46 -04:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Fix: safer and cleaner forward() in distributed embedding layer

Browse Source
This commit is contained in:
saro1993 2025-04-05 21:33:46 +02:00
parent a878eada08
commit abaadd9b3e
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:
ValueError: If `world_size` is not defined.
"""
if 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)
if world_size > 1:
y[mask] = 0
dist.all_reduce(y)
return y
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.world_size < 1:
raise ValueError("world_size must be >= 1")
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:
"""