zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-10-24 06:28:07 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix deepseek sharding (#1242)

Browse Source
This commit is contained in:
Awni Hannun
2025-02-03 16:59:50 -08:00
committed by GitHub
parent 0989c073b0
commit 21d0ab6e8a
3 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
llms/mlx_lm/models/deepseek_v2.py
View File

@@ -386,6 +386,7 @@ class DeepseekV2Model(nn.Module):
self.num_layers = layers_per_rank
self.layers = self.layers[: self.end_idx]
self.layers[: self.start_idx] = [None] * self.start_idx
self.num_layers = len(self.layers) - self.start_idx
def __call__(
self,

3
llms/mlx_lm/models/deepseek_v3.py
View File

@@ -397,12 +397,11 @@ class DeepseekV3Model(nn.Module):
layers_per_rank = (
len(self.layers) + self.pipeline_size - 1
) // self.pipeline_size
start = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
self.end_idx = self.start_idx + layers_per_rank
self.num_layers = layers_per_rank
self.layers = self.layers[: self.end_idx]
self.layers[: self.start_idx] = [None] * self.start_idx
self.num_layers = len(self.layers) - self.start_idx
def __call__(
self,