Deepseek moe #2467
Deepseek moe #2467
Conversation
| else: | ||
| final_hidden_states.add_(current_hidden_states) | ||
|
|
||
| y = tensor_model_parallel_all_reduce(final_hidden_states) |
There was a problem hiding this comment.
reducing after forwarding sharedexpert could avoiding one extra reduce operation, as sharedexpert also requires a reduce operation; these two operations can be merged.
There was a problem hiding this comment.
I didn't split shared_expert weights, and there is no need to reduce for that. Am i right?
There was a problem hiding this comment.
splitting shared_expert may have higher performance, reducing memory required by each gpu
There was a problem hiding this comment.
Emm..., It seems make sense. But maybe we can wait for #2293. And then consider improving performance.
| self.expert_indicies = np.array_split(range( | ||
| self.n_routed_experts), self.tp_size)[self.rank].tolist() |
There was a problem hiding this comment.
nit: we can replace with pure-pytorch version
torch.arange(self.n_routed_experts
).split(self.n_routed_experts // self.tp_size)[self.rank].tolist()
| @@ -0,0 +1,468 @@ | |||
| # coding=utf-8 | |||
| # Adapted from | |||
| # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py | |||
There was a problem hiding this comment.
can we update this comment? seems adapted mostly from mixtral.py
| max_position_embeddings=max_position_embeddings, | ||
| linear_method=linear_method, | ||
| ) | ||
| self.mlp = DeepseekMoE(config=config, |
There was a problem hiding this comment.
I'm pretty sure this would be easier to read if you did
if <...>:
self.mlp = DeepseekMoE(...)
else:
self.mlp = DeepseekMLP(...)
:)
|
@cadedaniel @pcmoritz Thanks for your reviews! I suggest to merge the fused moe version in #2453 which is faster. And i will close this PR and let's pushing that PR being merged. |
This is a refactor version based #2453, since i have no permission to change onto that PR.
load_weightlogic cleaner, and remove unnecessarymerged_replicated_linear_loaderfunction.reduce_resultsargs of DeepseekMLPshared_expertandexpertuse sameDeepseekExpertMLPwhich is semantically appropriate.MixtralMLP withDeepseekExpertMLPby usingReplicatedLinearacrossgate_proj,up_proj,down_proj.@zwd003 Thanks for your contribution again! And happy to be co-auther with you.
@zhuohan123 @WoosukKwon This PR is ready for review! Also we can wait for #2293