Commit 2743daa4 authored by zhuwenwen's avatar zhuwenwen
Browse files

update deepseek_v2.py

parent badc34b9
...@@ -257,9 +257,6 @@ class DeepseekV2MoE(nn.Module): ...@@ -257,9 +257,6 @@ class DeepseekV2MoE(nn.Module):
reduce_results=False, reduce_results=False,
prefix=f"{prefix}.shared_experts", prefix=f"{prefix}.shared_experts",
) )
from vllm.two_batch_overlap.two_batch_overlap import tbo_all_reduce
self.tbo_all_reduce = tbo_all_reduce
self.experts = SharedFusedMoE( self.experts = SharedFusedMoE(
shared_experts=self.shared_experts, shared_experts=self.shared_experts,
num_experts=config.n_routed_experts, num_experts=config.n_routed_experts,
...@@ -281,6 +278,9 @@ class DeepseekV2MoE(nn.Module): ...@@ -281,6 +278,9 @@ class DeepseekV2MoE(nn.Module):
num_redundant_experts=self.n_redundant_experts, num_redundant_experts=self.n_redundant_experts,
is_sequence_parallel=self.is_sequence_parallel, is_sequence_parallel=self.is_sequence_parallel,
) )
from vllm.two_batch_overlap.two_batch_overlap import tbo_all_reduce
self.tbo_all_reduce = tbo_all_reduce
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment