debug

modeling_jetmoe.py

@@ -654,7 +654,7 @@ class JetMoEAttention(nn.Module):
         self.num_heads = self.num_key_value_heads * self.top_k
         self.hidden_size_per_attention_head = config.kv_channels
 
-        self.experts = moe.MoE(
+        self.experts = MoE(
             input_size=config.hidden_size,
             hidden_size=self.kv_projection_size,
             num_experts=config.moe_num_experts,
@@ -1072,7 +1072,7 @@ class JetMoEBlock(nn.Module):
         # moe_args.activation_fn = F.silu
         # moe_args.return_bias = False
         # self.mlp = megablocks.layers.dmoe.dMoE(moe_args)
-        self.mlp = moe.MoE(
+        self.mlp = MoE(
             input_size=config.hidden_size,
             hidden_size=config.ffn_hidden_size,
             num_experts=config.moe_num_experts,